Verified Commit f2c67baa authored by Adel Kara Slimane's avatar Adel Kara Slimane
Browse files

Save initial onsite values for potential performance improvement.

The potential performance improvement may come from cythonised dict<int> lookup
vs pure pythong function call of the Hamiltonian value.

Further performance improvement can be made by making tinyarray cimport-able
and use libcpp.map instead of dict when saving static onsite values.
parent a0ce3c1a
Pipeline #43705 passed with stages
in 5 minutes and 39 seconds
......@@ -6,12 +6,13 @@ from kwant.graph.defs cimport gint
from kwant.operator cimport Current as ParticleCurrent
cdef class EnergyCurrent:
cdef public int sum, where_size, operatorType
cdef public int sum, where_size, operatorType, _saved_static_onsite
cdef public object syst, customOnsite
cdef public str time_name
cdef public float time_start
cdef public gint[:, :] where, _site_ranges
cdef public list _neighbors_i, _neighbors_j # List of lists of the neighbors of each site i and j in of the (j, i) hoppings in the `where` list
cdef public dict _static_onsite
cdef class EnergySource:
cdef public int sum, where_size, operatorType, has_dt_hamiltonian, has_dt_customOnsite, _saved_static_onsite
......@@ -24,12 +25,13 @@ cdef class EnergySource:
cdef public dict _static_onsite
cdef class EnergyDensity:
cdef public int sum, where_size, operatorType
cdef public int sum, where_size, operatorType, _saved_static_onsite
cdef public object syst, customOnsite
cdef public str time_name
cdef public float time_start
cdef public gint[:, :] _site_ranges, where
cdef public list _neighbors_i # List of lists of the neighbors of each site i in the `where` list
cdef public dict _static_onsite
cdef class LeadHeatCurrent:
cdef public object syst
......
......@@ -292,6 +292,8 @@ cdef class EnergyCurrent:
self.where_size = self.where.shape[0]
self._neighbors_i = list()
self._neighbors_j = list()
self._static_onsite = dict()
self._saved_static_onsite = 0
self._check_if_hopping_is_at_interface()
self._update_neighbor_lists()
......@@ -320,6 +322,20 @@ cdef class EnergyCurrent:
self._neighbors_i.append(np.asarray([k for k in self.syst.graph.out_neighbors(i)], dtype=gint_dtype))
self._neighbors_j.append(np.asarray([k for k in self.syst.graph.out_neighbors(j)], dtype=gint_dtype))
def _save_static_onsite(self, params=None):
init_params = params.copy()
init_params[self.time_name] = self.time_start - 1.0
for hopping_index in range(self.where_size):
i = self.where[hopping_index, 1]
j = self.where[hopping_index, 0]
if i not in self._static_onsite.keys():
self._static_onsite[i] = ta.matrix(self.syst.hamiltonian(i, i, params=init_params), complex)
if j not in self._static_onsite.keys():
self._static_onsite[j] = ta.matrix(self.syst.hamiltonian(j, j, params=init_params), complex)
@cython.boundscheck(False)
@cython.wraparound(False)
......@@ -373,6 +389,10 @@ cdef class EnergyCurrent:
if psi.shape != (tot_norbs,):
raise ValueError('psi vector is of incorrect shape')
if not self._saved_static_onsite:
self._saved_static_onsite = 1
self._save_static_onsite(params=params)
cdef complex[:] wf = psi
result = np.zeros(self.where_size, dtype=float)
......@@ -401,10 +421,8 @@ cdef class EnergyCurrent:
E_jj = ta.matrix(self.syst.hamiltonian(j, j, params=params), complex)
elif self.operatorType == 1:
# We get the hamiltonian at instant 0 as onsite
init_params = params.copy()
init_params[self.time_name] = self.time_start - 1.0
E_ii = ta.matrix(self.syst.hamiltonian(i, i, params=init_params), complex)
E_jj = ta.matrix(self.syst.hamiltonian(j, j, params=init_params), complex)
E_ii = self._static_onsite[i]
E_jj = self._static_onsite[j]
else:
# Use user provided onsite
E_ii = ta.matrix(self.customOnsite(self.syst.sites[i], params=params), complex)
......@@ -777,9 +795,7 @@ cdef class EnergySource:
ta.matrix(self.syst.hamiltonian(i, i, params=params_minus_dt), complex)) / (2*self.dt)
elif self.operatorType == 1:
# We get the hamiltonian at instant 0 as onsite
init_params = params.copy()
init_params[self.time_name] = self.time_start - 1.0
# We get the hamiltonian at instant 0 as onsite
V_i = ta.matrix(self.syst.hamiltonian(i, i, params=params), complex) - self._static_onsite[i]
# In the kinetic+ case the onsite has no time derivative
......@@ -831,9 +847,6 @@ cdef class EnergySource:
V_k = None
if self.operatorType == 1:
init_params = params.copy()
init_params[self.time_name] = self.time_start - 1.0
V_k = ta.matrix(self.syst.hamiltonian(k, k, params=params), complex) - self._static_onsite[k]
else:
V_k = ta.matrix(self.syst.hamiltonian(k, k, params=params), complex) -\
......@@ -1008,6 +1021,9 @@ cdef class EnergyDensity:
self.where_size = self.where.shape[0]
self._neighbors_i = list()
self._static_onsite = dict()
self._saved_static_onsite = 0
self._update_neighbor_lists()
self.time_name = _common.get_default_function_argument(onebody.WaveFunction.from_kwant, 'time_name')
......@@ -1027,8 +1043,16 @@ cdef class EnergyDensity:
for hopping in self.where:
i = hopping[0]
assert(isinstance(i, (int, np.integer)))
self._neighbors_i.append(np.asarray([k for k in self.syst.graph.out_neighbors(i)], dtype=gint_dtype))
self._neighbors_i.append(np.asarray([k for k in self.syst.graph.out_neighbors(i)], dtype=gint_dtype))
def _save_static_onsite(self, params=None):
init_params = params.copy()
init_params[self.time_name] = self.time_start - 1.0
for hopping_index in range(self.where_size):
i = self.where[hopping_index, 0]
if i not in self._static_onsite.keys():
self._static_onsite[i] = ta.matrix(self.syst.hamiltonian(i, i, params=init_params), complex)
@cython.boundscheck(False)
@cython.wraparound(False)
......@@ -1080,7 +1104,11 @@ cdef class EnergyDensity:
tot_norbs = _get_tot_norbs(self.syst)
if psi.shape != (tot_norbs,):
raise ValueError('psi vector is of incorrect shape')
raise ValueError('psi vector is of incorrect shape')
if not self._saved_static_onsite:
self._saved_static_onsite = 1
self._save_static_onsite(params=params)
cdef complex[:] wf = psi
......@@ -1106,10 +1134,7 @@ cdef class EnergyDensity:
elif self.operatorType == 1:
# We get the hamiltonian at instant 0 as onsite
init_params = params.copy()
init_params[self.time_name] = self.time_start - 1.0
E_ii = ta.matrix(self.syst.hamiltonian(i, i, params=init_params), complex)
E_ii = self._static_onsite[i]
else:
# Use user provided onsite
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment