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
......
Markdown is supported
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