operator.pyx 58 KB
 Joseph Weston committed Dec 10, 2019 1 # Copyright 2011-2019 Kwant authors.  Joseph Weston committed Aug 04, 2016 2 3 4 # # This file is part of Kwant. It is subject to the license terms in the file # LICENSE.rst found in the top-level directory of this distribution and at  Anton Akhmerov committed Dec 22, 2020 5 # https://kwant-project.org/license. A list of Kwant authors can be found in  Joseph Weston committed Aug 04, 2016 6 # the file AUTHORS.rst at the top-level directory of this distribution and at  Anton Akhmerov committed Dec 22, 2020 7 # https://kwant-project.org/authors.  Joseph Weston committed Dec 15, 2016 8 """Tools for working with operators for acting on wavefunctions."""  Joseph Weston committed Aug 04, 2016 9 10 11 12  __all__ = ['Density', 'Current', 'Source'] import cython  Joseph Weston committed Dec 10, 2019 13 import itertools  Joseph Weston committed Aug 04, 2016 14 15 import functools as ft import collections  Joseph Weston committed Nov 13, 2019 16 import warnings  Anton Akhmerov committed Jan 22, 2017 17   Joseph Weston committed Aug 04, 2016 18 19 import numpy as np import tinyarray as ta  Anton Akhmerov committed Jan 22, 2017 20 from scipy.sparse import coo_matrix  Joseph Weston committed Aug 04, 2016 21 22 23  from libc cimport math  Joseph Weston committed Dec 10, 2019 24 25 26 cdef extern from "complex.h": double cabs(double complex)  Joseph Weston committed Aug 04, 2016 27 from .graph.core cimport EdgeIterator  Joseph Weston committed Aug 28, 2019 28 from .graph.core import DisabledFeatureError, NodeDoesNotExistError  Joseph Weston committed Aug 04, 2016 29 30 from .graph.defs cimport gint from .graph.defs import gint_dtype  Joseph Weston committed Dec 10, 2019 31 32 33 from .system import ( is_infinite, is_vectorized, Site, SiteArray, _normalize_matrix_blocks )  Joseph Weston committed Nov 13, 2019 34 35 36 from ._common import ( UserCodeError, KwantDeprecationWarning, get_parameters, deprecate_args )  Joseph Weston committed Aug 04, 2016 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56  ################ Generic Utility functions @cython.boundscheck(False) @cython.wraparound(False) cdef gint _bisect(gint[:] a, gint x): "bisect.bisect specialized for searching site_ranges" cdef gint mid, lo = 0, hi = a.shape[0] while lo < hi: mid = (lo + hi) // 2 if x < a[mid]: hi = mid else: lo = mid + 1 return lo @cython.boundscheck(False) @cython.wraparound(False)  Joseph Weston committed Dec 10, 2019 57 58 59 60 61 62 63 cdef int _is_hermitian( complex[:, :] a, double atol=1e-300, double rtol=1e-13 ) except -1: "Return True if 'a' is Hermitian" if a.shape[0] != a.shape[1]: return False  Joseph Weston committed Aug 04, 2016 64 65 66  # compute max(a) cdef double tmp, max_a = 0  Joseph Weston committed Dec 10, 2019 67  cdef gint i, j, k  Joseph Weston committed Aug 04, 2016 68 69  for i in range(a.shape[0]): for j in range(a.shape[1]):  Joseph Weston committed Dec 10, 2019 70  tmp = cabs(a[i, j])  Joseph Weston committed Aug 04, 2016 71 72 73 74 75 76  if tmp > max_a: max_a = tmp max_a = math.sqrt(max_a) cdef double tol = rtol * max_a + atol for i in range(a.shape[0]):  Joseph Weston committed Dec 10, 2019 77 78  for j in range(i, a.shape[1]): tmp = cabs(a[i, j] - a[j, i].conjugate())  Joseph Weston committed Aug 04, 2016 79 80 81 82  if tmp > tol: return False return True  Joseph Weston committed Dec 10, 2019 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 @cython.boundscheck(False) @cython.wraparound(False) cdef int _is_hermitian_3d( complex[:, :, :] a, double atol=1e-300, double rtol=1e-13 ) except -1: "Return True if 'a' is Hermitian" if a.shape[1] != a.shape[2]: return False # compute max(a) cdef double tmp, max_a = 0 cdef gint i, j, k for k in range(a.shape[0]): for i in range(a.shape[1]): for j in range(a.shape[2]): tmp = cabs(a[k, i, j]) if tmp > max_a: max_a = tmp max_a = math.sqrt(max_a) cdef double tol = rtol * max_a + atol for k in range(a.shape[0]): for i in range(a.shape[1]): for j in range(i, a.shape[2]): tmp = cabs(a[k, i, j] - a[k, j, i].conjugate()) if tmp > tol: return False return True  Joseph Weston committed Aug 04, 2016 112   Joseph Weston committed Dec 10, 2019 113 114 115 116 117 118 119 120 121 122 123 124 125 126  @cython.boundscheck(False) @cython.wraparound(False) cdef _select(gint[:, :] arr, gint[:] indexes): ret = np.empty((indexes.shape[0], arr.shape[1]), dtype=gint_dtype) cdef gint[:, :] ret_view = ret cdef gint i, j for i in range(indexes.shape[0]): for j in range(arr.shape[1]): ret_view[i, j] = arr[indexes[i], j] return ret  Joseph Weston committed Aug 04, 2016 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 ################ Helper functions _shape_msg = ('{0} matrix dimensions do not match ' 'the declared number of orbitals') _herm_msg = ('{0} matrix is not hermitian, use the option ' 'check_hermiticity=True if this is intentional.') cdef int _check_onsite(complex[:, :] M, gint norbs, int check_hermiticity) except -1: "Check onsite matrix for correct shape and hermiticity." if M.shape[0] != M.shape[1]: raise UserCodeError('Onsite matrix is not square') if M.shape[0] != norbs: raise UserCodeError(_shape_msg.format('Onsite'))  Joseph Weston committed Dec 10, 2019 142  if check_hermiticity and not _is_hermitian(M):  Joseph Weston committed Aug 04, 2016 143 144 145 146  raise ValueError(_herm_msg.format('Onsite')) return 0  Joseph Weston committed Dec 10, 2019 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 cdef int _check_onsites(complex[:, :, :] M, gint norbs, int check_hermiticity) except -1: "Check onsite matrix for correct shape and hermiticity." if M.shape[1] != M.shape[2]: raise UserCodeError('Onsite matrix is not square') if M.shape[1] != norbs: raise UserCodeError(_shape_msg.format('Onsite')) if check_hermiticity and not _is_hermitian_3d(M): raise ValueError(_herm_msg.format('Onsite')) return 0 cdef int _check_hams(complex[:, :, :] H, gint to_norbs, gint from_norbs, int check_hermiticity) except -1: if H.shape[1] != to_norbs or H.shape[2] != from_norbs:  Joseph Weston committed Aug 04, 2016 162  raise UserCodeError(_shape_msg.format('Hamiltonian'))  Joseph Weston committed Dec 10, 2019 163  if check_hermiticity and not _is_hermitian_3d(H):  Joseph Weston committed Aug 04, 2016 164 165 166 167  raise ValueError(_herm_msg.format('Hamiltonian')) return 0  Pablo Piskunow committed Dec 25, 2020 168 169 170 171 172 173 174 175 176 177 def _check_norbs(syst): # TODO: Update this when it becomes clear how ND systems will be # implemented. if syst.site_ranges is None: raise ValueError('Number of orbitals not defined.\n' 'Declare the number of orbitals using the ' 'norbs keyword argument when constructing ' 'the site families (lattices).')  Joseph Weston committed Aug 04, 2016 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 @cython.boundscheck(False) @cython.wraparound(False) cdef void _get_orbs(gint[:, :] site_ranges, gint site, gint *start_orb, gint *norbs): """Return the first orbital of this site and the number of orbitals""" cdef gint run_idx, first_site, norb, orb_offset, orb # Calculate the index of the range that contains the site. run_idx = _bisect(site_ranges[:, 0], site) - 1 first_site = site_ranges[run_idx, 0] norb = site_ranges[run_idx, 1] orb_offset = site_ranges[run_idx, 2] # calculate the slice start_orb[0] = orb_offset + (site - first_site) * norb norbs[0] = norb @cython.boundscheck(False) @cython.wraparound(False) def _get_all_orbs(gint[:, :] where, gint[:, :] site_ranges): cdef gint[:, :] offsets = np.empty((where.shape[0], 2), dtype=gint_dtype) cdef gint[:, :] norbs = np.empty((where.shape[0], 2), dtype=gint_dtype) cdef gint w, a, a_offset, a_norbs, b, b_offset, b_norbs for w in range(where.shape[0]): a = where[w, 0] _get_orbs(site_ranges, a, &a_offset, &a_norbs) if where.shape[1] == 1: b, b_offset, b_norbs = a, a_offset, a_norbs else: b = where[w, 1] _get_orbs(site_ranges, b, &b_offset, &b_norbs) offsets[w, 0] = a_offset offsets[w, 1] = b_offset norbs[w, 0] = a_norbs norbs[w, 1] = b_norbs return offsets, norbs def _get_tot_norbs(syst): cdef gint _unused, tot_norbs  Pablo Piskunow committed Dec 25, 2020 219  _check_norbs(syst)  Joseph Weston committed Nov 13, 2019 220  is_infinite_system = is_infinite(syst)  Joseph Weston committed Aug 04, 2016 221 222 223 224 225 226 227 228 229 230  n_sites = syst.cell_size if is_infinite_system else syst.graph.num_nodes _get_orbs(np.asarray(syst.site_ranges, dtype=gint_dtype), n_sites, &tot_norbs, &_unused) return tot_norbs def _normalize_site_where(syst, where): """Normalize the format of where when where contains sites. If where is None, then all sites in the system are returned.  Joseph Weston committed Aug 28, 2019 231 232  If it is a general sequence then it is expanded into an array. If syst is a finalized Builder then where may contain Site objects,  Joseph Weston committed Aug 04, 2016 233 234 235  otherwise it should contain integers. """ if where is None:  Joseph Weston committed Nov 13, 2019 236  if is_infinite(syst):  Joseph Weston committed Aug 28, 2019 237 238 239  where = list(range(syst.cell_size)) else: where = list(range(syst.graph.num_nodes))  Joseph Weston committed Aug 04, 2016 240  elif callable(where):  Anton Akhmerov committed Jan 24, 2017 241  try:  Joseph Weston committed Aug 28, 2019 242  where = [syst.id_by_site[s] for s in filter(where, syst.sites)]  Anton Akhmerov committed Jan 24, 2017 243  except AttributeError:  Joseph Weston committed Nov 13, 2019 244  if is_infinite(syst):  Joseph Weston committed Aug 28, 2019 245 246 247  where = [s for s in range(syst.cell_size) if where(s)] else: where = [s for s in range(syst.graph.num_nodes) if where(s)]  Joseph Weston committed Aug 04, 2016 248  else:  Joseph Weston committed Nov 13, 2019 249  if isinstance(where[0], Site):  Joseph Weston committed Aug 28, 2019 250 251 252 253 254  try: where = [syst.id_by_site[s] for s in where] except AttributeError: raise TypeError("'where' contains Sites, but the system is not " "a finalized Builder.")  Joseph Weston committed Aug 04, 2016 255   Joseph Weston committed Aug 28, 2019 256  where = np.asarray(where, dtype=gint_dtype).reshape(-1, 1)  Joseph Weston committed Aug 04, 2016 257   Joseph Weston committed Nov 13, 2019 258  if is_infinite(syst) and np.any(where >= syst.cell_size):  Joseph Weston committed Aug 28, 2019 259 260 261 262 263 264 265  raise ValueError('Only sites in the fundamental domain may be ' 'specified using where.') if np.any(np.logical_or(where < 0, where >= syst.graph.num_nodes)): raise ValueError('where contains sites that are not in the ' 'system.') return where  Joseph Weston committed Aug 04, 2016 266 267 268 269 270 271 272  def _normalize_hopping_where(syst, where): """Normalize the format of where when where contains hoppings. If where is None, then all hoppings in the system are returned. If it is a general iterator then it is expanded into an array. If syst is  Joseph Weston committed Aug 28, 2019 273  a finalized Builder then where may contain pairs of Site objects,  Joseph Weston committed Aug 04, 2016 274 275 276 277 278  otherwise it should contain pairs of integers. """ if where is None: # we cannot extract the hoppings in the same order as they are in the # graph while simultaneously excluding all inter-cell hoppings  Joseph Weston committed Nov 13, 2019 279  if is_infinite(syst):  Joseph Weston committed Aug 04, 2016 280 281  raise ValueError('where must be provided when calculating ' 'current in an InfiniteSystem.')  Joseph Weston committed Aug 28, 2019 282  where = list(syst.graph)  Joseph Weston committed Aug 04, 2016 283  elif callable(where):  Anton Akhmerov committed Jan 24, 2017 284  if hasattr(syst, "sites"):  Joseph Weston committed Aug 28, 2019 285  def idxwhere(hop):  Joseph Weston committed Aug 04, 2016 286 287  a, b = hop return where(syst.sites[a], syst.sites[b])  Joseph Weston committed Aug 28, 2019 288  where = list(filter(idxwhere, syst.graph))  Joseph Weston committed Aug 04, 2016 289  else:  Joseph Weston committed Aug 28, 2019 290  where = list(filter(lambda h: where(*h), syst.graph))  Joseph Weston committed Aug 04, 2016 291  else:  Joseph Weston committed Nov 13, 2019 292  if isinstance(where[0][0], Site):  Joseph Weston committed Aug 28, 2019 293 294 295 296 297 298 299 300  try: where = list((syst.id_by_site[a], syst.id_by_site[b]) for a, b in where) except AttributeError: raise TypeError("'where' contains Sites, but the system is not " "a finalized Builder.") # NOTE: if we ever have operators that contain elements that are # not in the system graph, then we should modify this check  Anton Akhmerov committed Jan 24, 2017 301  try:  Joseph Weston committed Aug 28, 2019 302 303  error = ValueError('where contains hoppings that are not ' 'in the system.')  Joseph Weston committed Aug 04, 2016 304  if any(not syst.graph.has_edge(*w) for w in where):  Joseph Weston committed Aug 28, 2019 305 306 307 308 309 310  raise error # If where contains: negative integers, or integers > # of sites except (NodeDoesNotExistError, DisabledFeatureError): raise error where = np.asarray(where, dtype=gint_dtype)  Joseph Weston committed Aug 04, 2016 311   Joseph Weston committed Nov 13, 2019 312  if is_infinite(syst) and np.any(where > syst.cell_size):  Joseph Weston committed Aug 28, 2019 313 314  raise ValueError('Only intra-cell hoppings may be specified ' 'using where.')  Joseph Weston committed Aug 04, 2016 315   Joseph Weston committed Aug 28, 2019 316  return where  Joseph Weston committed Aug 04, 2016 317 318   Joseph Weston committed Dec 10, 2019 319 ## These four classes are here to avoid using closures, as these will  Joseph Weston committed Mar 17, 2017 320 321 ## break pickle support. These are only used inside '_normalize_onsite'.  Joseph Weston committed Dec 10, 2019 322 323 324 325 326 327 328 329 330 331 def _raise_user_error(exc, func, vectorized): msg = [ 'Error occurred in user-supplied onsite function "{0}".', 'Did you remember to vectorize "{0}"?' if vectorized else '', 'See the upper part of the above backtrace for more information.', ] msg = '\n'.join(line for line in msg if line).format(func.__name__) raise UserCodeError(msg) from exc  Joseph Weston committed Mar 17, 2017 332 333 class _FunctionalOnsite:  Joseph Weston committed Dec 10, 2019 334  def __init__(self, onsite, sites, site_ranges):  Joseph Weston committed Mar 17, 2017 335 336  self.onsite = onsite self.sites = sites  Joseph Weston committed Dec 10, 2019 337 338 339 340  self.site_ranges = site_ranges def __call__(self, site_range, site_offsets, *args): sites = self.sites  Joseph Weston committed Dec 12, 2019 341  offset, norbs, _ = self.site_ranges[site_range]  Joseph Weston committed Dec 10, 2019 342 343 344 345  try: ret = [self.onsite(sites[offset + i], *args) for i in site_offsets] except Exception as exc: _raise_user_error(exc, self.onsite, vectorized=False)  Joseph Weston committed Dec 12, 2019 346 347  expected_shape = (len(site_offsets), norbs, norbs)  Joseph Weston committed Dec 12, 2019 348 349  return _normalize_matrix_blocks(ret, expected_shape, calling_function=self.onsite)  Joseph Weston committed Dec 10, 2019 350   Joseph Weston committed Mar 17, 2017 351   Joseph Weston committed Dec 10, 2019 352 353 354 355 356 357 358 359 360 361 362 363 364 365 class _VectorizedFunctionalOnsite: def __init__(self, onsite, site_arrays): self.onsite = onsite self.site_arrays = site_arrays def __call__(self, site_range, site_offsets, *args): site_array = self.site_arrays[site_range] tags = site_array.tags[site_offsets] sites = SiteArray(site_array.family, tags) try: ret = self.onsite(sites, *args) except Exception as exc: _raise_user_error(exc, self.onsite, vectorized=True)  Joseph Weston committed Dec 12, 2019 366 367  expected_shape = (len(sites), sites.family.norbs, sites.family.norbs)  Joseph Weston committed Dec 12, 2019 368 369  return _normalize_matrix_blocks(ret, expected_shape, calling_function=self.onsite)  Joseph Weston committed Mar 17, 2017 370 371   Joseph Weston committed Dec 10, 2019 372 class _FunctionalOnsiteNoTransform:  Joseph Weston committed Mar 17, 2017 373   Joseph Weston committed Dec 10, 2019 374 375 376 377 378  def __init__(self, onsite, site_ranges): self.onsite = onsite self.site_ranges = site_ranges def __call__(self, site_range, site_offsets, *args):  Joseph Weston committed Dec 12, 2019 379 380  offset, norbs, _ = self.site_ranges[site_range] site_ids = offset + site_offsets  Joseph Weston committed Dec 10, 2019 381 382 383 384  try: ret = [self.onsite(id, *args) for id in site_ids] except Exception as exc: _raise_user_error(exc, self.onsite, vectorized=False)  Joseph Weston committed Dec 12, 2019 385 386  expected_shape = (len(site_offsets), norbs, norbs)  Joseph Weston committed Dec 12, 2019 387 388  return _normalize_matrix_blocks(ret, expected_shape, calling_function=self.onsite)  Joseph Weston committed Dec 10, 2019 389 390 391 392 393 394 395 396 397 398 399  class _DictOnsite: def __init__(self, onsite, range_family_map): self.onsite = onsite self.range_family_map = range_family_map def __call__(self, site_range, site_offsets, *args): fam = self.range_family_map[site_range] ret = [self.onsite[fam]] * len(site_offsets)  Joseph Weston committed Dec 12, 2019 400 401 402  expected_shape = (len(site_offsets), fam.norbs, fam.norbs) return _normalize_matrix_blocks(ret, expected_shape)  Joseph Weston committed Mar 17, 2017 403 404   Joseph Weston committed Aug 04, 2016 405 406 407 408 409 410 def _normalize_onsite(syst, onsite, check_hermiticity): """Normalize the format of onsite. If onsite is a function or a mapping (dictionary) then a function is returned. """  Christoph Groth committed Sep 20, 2018 411  param_names = ()  Joseph Weston committed Feb 21, 2017 412   Joseph Weston committed Aug 04, 2016 413  if callable(onsite):  Joseph Weston committed Feb 21, 2017 414  # make 'onsite' compatible with hamiltonian value functions  Christoph Groth committed Sep 20, 2018 415  param_names = get_parameters(onsite)[1:]  Joseph Weston committed Dec 10, 2019 416 417 418 419 420 421  if is_vectorized(syst): _onsite = _VectorizedFunctionalOnsite(onsite, syst.site_arrays) elif hasattr(syst, "sites"): # probably a non-vectorized finalized Builder _onsite = _FunctionalOnsite(onsite, syst.sites, syst.site_ranges) else: # generic old-style system, therefore *not* vectorized. _onsite = _FunctionalOnsiteNoTransform(onsite, syst.site_ranges)  Joseph Weston committed Aug 04, 2016 422   Joseph Weston committed Dec 10, 2019 423  elif isinstance(onsite, collections.abc.Mapping):  Joseph Weston committed Aug 04, 2016 424 425 426 427 428  # onsites known; immediately check for correct shape and hermiticity for fam, _onsite in onsite.items(): _onsite = ta.matrix(_onsite, complex) _check_onsite(_onsite, fam.norbs, check_hermiticity)  Joseph Weston committed Dec 10, 2019 429 430 431 432 433 434 435 436 437 438  if is_vectorized(syst): range_family_map = [arr.family for arr in syst.site_arrays] elif not hasattr(syst, 'sites'): raise TypeError('Provide onsite as a value or a function for ' 'systems that are not finalized Builders.') else: # The last entry in 'site_ranges' is just an end marker, so we remove it range_family_map = [syst.sites[r[0]].family for r in syst.site_ranges[:-1]] _onsite = _DictOnsite(onsite, range_family_map)  Joseph Weston committed Aug 04, 2016 439 440 441 442 443 444 445 446 447 448  else: # single onsite; immediately check for correct shape and hermiticity _onsite = ta.matrix(onsite, complex) _check_onsite(_onsite, _onsite.shape[0], check_hermiticity) if _onsite.shape[0] == 1: # NOTE: this is wasteful when many orbitals per site, but it # simplifies the code in _operate. If this proves to be a # bottleneck, then we can add a code path for scalar onsites max_norbs = max(norbs for (_, norbs, _) in syst.site_ranges) _onsite = _onsite[0, 0] * ta.identity(max_norbs, complex)  Joseph Weston committed Dec 06, 2019 449  elif len(set(norbs for _, norbs, _ in syst.site_ranges[:-1])) == 1:  Joseph Weston committed Aug 04, 2016 450 451 452 453 454 455 456 457 458 459 460  # we have the same number of orbitals everywhere norbs = syst.site_ranges[0][1] if _onsite.shape[0] != norbs: msg = ('Single onsite matrix of shape ({0}, {0}) provided ' 'but there are {1} orbitals per site in the system') raise ValueError(msg.format(_onsite.shape[0], norbs)) else: msg = ('Single onsite matrix provided, but there are ' 'different numbers of orbitals on different sites') raise ValueError(msg)  Christoph Groth committed Sep 20, 2018 461  return _onsite, param_names  Joseph Weston committed Aug 04, 2016 462 463   Joseph Weston committed Dec 10, 2019 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 def _make_onsite_or_hopping_terms(site_ranges, where): terms = dict() cdef gint[:] site_offsets_ = np.asarray(site_ranges, dtype=gint_dtype)[:, 0] cdef gint i if where.shape[1] == 1: # onsite for i in range(where.shape[0]): a = _bisect(site_offsets_, where[i, 0]) - 1 terms.setdefault((a, a), []).append(i) else: # hopping for i in range(where.shape[0]): a = _bisect(site_offsets_, where[i, 0]) - 1 b = _bisect(site_offsets_, where[i, 1]) - 1 terms.setdefault((a, b), []).append(i) return [(a, None, b) for a, b in terms.items()] def _vectorized_make_onsite_terms(syst, where): assert is_vectorized(syst) assert where.shape[1] == 1 site_offsets = [r[0] for r in syst.site_ranges] terms = {} term_by_id = syst._onsite_term_by_site_id for i in range(where.shape[0]): w = where[i, 0] terms.setdefault(term_by_id[w], []).append(i) ret = [] for term_id, which in terms.items(): term = syst.terms[term_id]  Joseph Weston committed Dec 10, 2019 497 498  ((term_sa, _), (term_site_offsets, _)) = syst.subgraphs[term.subgraph] term_sites = site_offsets[term_sa] + term_site_offsets  Joseph Weston committed Dec 10, 2019 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558  which = np.asarray(which, dtype=gint_dtype) sites = _select(where, which).reshape(-1) selector = np.searchsorted(term_sites, sites) ret.append((term_id, selector, which)) return ret def _vectorized_make_hopping_terms(syst, where): assert is_vectorized(syst) assert where.shape[1] == 2 site_offsets = [r[0] for r in syst.site_ranges] terms = {} term_by_id = syst._hopping_term_by_edge_id for i in range(where.shape[0]): a, b = where[i, 0], where[i, 1] edge = syst.graph.first_edge_id(a, b) terms.setdefault(term_by_id[edge], []).append(i) ret = [] dtype = np.dtype([('f0', int), ('f1', int)]) for term_id, which in terms.items(): herm_conj = term_id < 0 if herm_conj: real_term_id = -term_id - 1 else: real_term_id = term_id which = np.asarray(which, dtype=gint_dtype) # Select out the hoppings and reverse them if we are # dealing with Hermitian conjugate hoppings hops = _select(where, which) if herm_conj: hops = hops[:, ::-1] # Force contiguous array to handle herm conj case also. # Needs to be contiguous to cast to compound dtype hops = np.ascontiguousarray(hops, dtype=int) hops = hops.view(dtype).reshape(-1) term = syst.terms[real_term_id] # Get array of pairs ((to_sa, from_sa), term_hoppings) = syst.subgraphs[term.subgraph] term_hoppings = term_hoppings.transpose() + (site_offsets[to_sa], site_offsets[from_sa]) term_hoppings = term_hoppings.view(dtype).reshape(-1) selector = np.recarray.searchsorted(term_hoppings, hops) ret.append((term_id, selector, which)) return ret def _make_matrix_elements(evaluate_term, terms): matrix_elements = [] for (term_id, term_selector, which) in terms: which = np.asarray(which, dtype=gint_dtype) data = evaluate_term(term_id, term_selector, which) matrix_elements.append((which, data)) return matrix_elements  Joseph Weston committed Aug 04, 2016 559 560 561 562 563 564 565 566 567 568 569 570 571 572 cdef class BlockSparseMatrix: """A sparse matrix stored as dense blocks. Parameters ---------- where : gint[:, :] Nx2 matrix or Nx1 matrix: the arguments a and b to be used when evaluating f. If an Nx1 matrix, then b=a. block_offsets : gint[:, :] The row and column offsets for the start of each block in the sparse matrix: (row_offset, col_offset). block_shapes : gint[:, :] Nx2 array: the shapes of each block, (n_rows, n_cols).  Joseph Weston committed Dec 10, 2019 573 574 575 576  matrix_elements : sequence of pairs (where_indices, data) 'data' is a 3D complex array; a vector of matrices. 'where_indices' is a 1D array of indices for 'where'; 'data[i]' should be placed at 'where[where_indices[i]]'.  Joseph Weston committed Aug 04, 2016 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594  Attributes ---------- block_offsets : gint[:, :] The row and column offsets for the start of each block in the sparse matrix: (row_offset, col_offset). block_shapes : gint[:, :] The shape of each block: (n_rows, n_cols) data_offsets : gint[:] The offsets of the start of each matrix block in data. data : complex[:] The matrix of each block, stored in row-major (C) order. """ @cython.embedsignature @cython.boundscheck(False) @cython.wraparound(False) def __init__(self, gint[:, :] where, gint[:, :] block_offsets,  Joseph Weston committed Dec 10, 2019 595  gint[:, :] block_shapes, matrix_elements):  Joseph Weston committed Aug 04, 2016 596 597 598 599 600 601 602 603 604 605 606 607 608 609  if (block_offsets.shape[0] != where.shape[0] or block_shapes.shape[0] != where.shape[0]): raise ValueError('Arrays should be the same length along ' 'the first axis.') self.block_shapes = block_shapes self.block_offsets = block_offsets self.data_offsets = np.empty(where.shape[0], dtype=gint_dtype) ### calculate shapes and data_offsets cdef gint w, data_size = 0 for w in range(where.shape[0]): self.data_offsets[w] = data_size data_size += block_shapes[w, 0] * block_shapes[w, 1] ### Populate data array self.data = np.empty((data_size,), dtype=complex)  Joseph Weston committed Dec 10, 2019 610 611 612 613 614 615 616 617 618 619 620 621 622  cdef complex[:, :, :] data cdef gint[:] where_indexes cdef gint i, j, k, off, a, b, a_norbs, b_norbs for where_indexes, data in matrix_elements: for i in range(where_indexes.shape[0]): w = where_indexes[i] off = self.data_offsets[w] a_norbs = self.block_shapes[w, 0] b_norbs = self.block_shapes[w, 1] # Copy data for j in range(a_norbs): for k in range(b_norbs): self.data[off + j * b_norbs + k] = data[i, j, k]  Joseph Weston committed Aug 04, 2016 623 624 625 626  cdef complex* get(self, gint block_idx): return &self.data[0] + self.data_offsets[block_idx]  Joseph Weston committed Mar 17, 2017 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641  def __getstate__(self): return tuple(map(np.asarray, ( self.block_offsets, self.block_shapes, self.data_offsets, self.data ))) def __setstate__(self, state): (self.block_offsets, self.block_shapes, self.data_offsets, self.data, ) = state  Joseph Weston committed Aug 04, 2016 642 643 644 645 646 647 648 649 650 651  ################ Local Observables # supported operations within the _operate method ctypedef enum operation: MAT_ELS ACT cdef class _LocalOperator:  Christoph Groth committed Feb 17, 2019 652 653  """Base class for operators defined by an on-site matrix and the Hamiltonian.  Joseph Weston committed Aug 04, 2016 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669  This includes "true" local operators, as well as "currents" and "sources". Attributes ---------- syst : ~kwant.system.System The system for which this operator is defined. Must have the number of orbitals defined for all site families. onsite : complex 2D array, or callable If a complex array, then the same onsite is used everywhere. Otherwise, function that can be called with a single site (integer) and extra arguments, and returns the representation of the operator on that site. This should return either a scalar or a square matrix of the same shape as that returned by the system Hamiltonian evaluated on the same site. The extra arguments must be the same as the extra arguments to syst.hamiltonian.  Joseph Weston committed Dec 15, 2016 670 671 672 673  where : 2D array of int or None where to evaluate the operator. A list of sites for on-site operators (accessed like where[n, 0]), otherwise a list of pairs of sites (accessed like where[n, 0] and where[n, 1]).  Joseph Weston committed Aug 04, 2016 674 675 676  check_hermiticity : bool If True, checks that onsite, as well as any relevant parts of the Hamiltonian are hermitian.  Joseph Weston committed Feb 02, 2017 677  sum : bool, default: False  Joseph Weston committed Dec 08, 2017 678 679 680  If True, then calling this operator will return a single scalar, otherwise a vector will be returned (see ~kwant.operator._LocalOperator.__call__ for details).  Joseph Weston committed Aug 04, 2016 681 682 683  """ @cython.embedsignature  Joseph Weston committed Dec 15, 2016 684 685  def __init__(self, syst, onsite, where, *, check_hermiticity=True, sum=False):  Pablo Piskunow committed Dec 25, 2020 686  _check_norbs(syst)  Joseph Weston committed Dec 10, 2019 687 688 689  if is_vectorized(syst) and is_infinite(syst): raise TypeError('Vectorized infinite systems cannot yet be ' 'used with operators.')  Joseph Weston committed Aug 04, 2016 690 691  self.syst = syst  Christoph Groth committed Sep 20, 2018 692 693  self.onsite, self._onsite_param_names = _normalize_onsite( syst, onsite, check_hermiticity)  Joseph Weston committed Aug 04, 2016 694  self.check_hermiticity = check_hermiticity  Joseph Weston committed Dec 15, 2016 695  self.sum = sum  Joseph Weston committed Aug 04, 2016 696  self._site_ranges = np.asarray(syst.site_ranges, dtype=gint_dtype)  Joseph Weston committed Dec 15, 2016 697  self.where = where  Joseph Weston committed Aug 04, 2016 698 699 700  self._bound_onsite = None self._bound_hamiltonian = None  Joseph Weston committed Dec 10, 2019 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734  # Here we pre-compute the datastructures that will enable us to evaluate # the Hamiltonian and onsite functions in a vectorized way. Essentially # we group the sites/hoppings in 'where' by what term of 'syst' they are # in (for vectorized systems), or by the site family(s) (for # non-vectorized systems). If the system is vectorized we store a list # of triples: # # (term_id, term_selector, which) # # otherwise # # ((to_site_range, from_site_range), None, which) # # Where: # # 'term_id' → integer: term index, may be negative (indicates herm. conj.) # 'term_selector' → 1D integer array: selects which elements from the # subgraph of term number 'term_id' should be evaluated. # 'which' → 1D integer array: selects which elements of 'where' this # vectorized evaluation corresponds to. # 'to/from_site_range' → integer: the site ranges that the elements of # 'where' indexed by 'which' correspond to. # # Note that all sites/hoppings indicated by 'which' belong to the *same* # pair of site families by construction. This is what allows for # vectorized evaluation. if is_vectorized(syst): if self.where.shape[1] == 1: self._terms = _vectorized_make_onsite_terms(syst, where) else: self._terms = _vectorized_make_hopping_terms(syst, where) else: self._terms = _make_onsite_or_hopping_terms(self._site_ranges, where)  Joseph Weston committed Aug 04, 2016 735  @cython.embedsignature  Joseph Weston committed Feb 21, 2017 736 737  def __call__(self, bra, ket=None, args=(), *, params=None): r"""Return the matrix elements of the operator.  Joseph Weston committed Aug 04, 2016 738   Joseph Weston committed Feb 02, 2017 739 740 741 742 743 744 745 746 747  An operator A can be called like >>> A(psi) to compute the expectation value :math:\bra{ψ} A \ket{ψ}, or like >>> A(phi, psi)  Joseph Weston committed Dec 08, 2017 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768  to compute the matrix element :math:\bra{φ} A \ket{ψ}. If sum=True was provided when constructing the operator, then a scalar is returned. If sum=False, then a vector is returned. The vector is defined over the sites of the system if the operator is a ~kwant.operator.Density, or over the hoppings if it is a ~kwant.operator.Current or ~kwant.operator.Source. By default, the returned vector is ordered in the same way as the sites (for ~kwant.operator.Density) or hoppings in the graph of the system (for ~kwant.operator.Current or ~kwant.operator.Density). If the keyword parameter where was provided when constructing the operator, then the returned vector is instead defined only over the sites or hoppings specified, and is ordered in the same way as where. Alternatively stated, for an operator :math:Q_{iαβ}, bra :math:φ_α and ket :math:ψ_β this computes :math:q_i = ∑_{αβ} φ^*_α Q_{iαβ} ψ_β if self.sum is False, otherwise computes :math:q = ∑_{iαβ} φ^*_α Q_{iαβ} ψ_β. where :math:i runs over all sites or hoppings, and :math:α and :math:β run over all the degrees of freedom.  Joseph Weston committed Dec 15, 2016 769   Joseph Weston committed Aug 04, 2016 770 771  Parameters ----------  Joseph Weston committed Dec 15, 2016 772  bra, ket : sequence of complex  Joseph Weston committed Aug 04, 2016 773 774 775 776 777 778  Must have the same length as the number of orbitals in the system. If only one is provided, both bra and ket are taken as equal. args : tuple, optional The arguments to pass to the system. Used to evaluate the onsite elements and, possibly, the system Hamiltonian.  Joseph Weston committed Feb 15, 2019 779  Deprecated in favor of 'params' (and mutually exclusive with it).  Joseph Weston committed Feb 21, 2017 780 781 782  params : dict, optional Dictionary of parameter names and their values. Mutually exclusive with 'args'.  Joseph Weston committed Aug 04, 2016 783 784 785  Returns -------  Joseph Weston committed Feb 02, 2017 786 787 788  float if check_hermiticity is True, and ket is None, otherwise complex. If this operator was created with sum=True, then a single value is returned, otherwise an array is returned.  Joseph Weston committed Aug 04, 2016 789  """  Joseph Weston committed Feb 21, 2017 790 791 792 793  if (self._bound_onsite or self._bound_hamiltonian) and (args or params): raise ValueError("Extra arguments are already bound to this " "operator. You should call this operator " "providing neither 'args' nor 'params'.")  Joseph Weston committed Feb 15, 2019 794 795 796 797 798  if args: # deprecate_args does not play nicely with methods of cdef classes, # when used as a decorator, so we manually raise the # deprecation warning here. deprecate_args()  Joseph Weston committed Feb 21, 2017 799 800  if args and params: raise TypeError("'args' and 'params' are mutually exclusive.")  Joseph Weston committed Aug 04, 2016 801 802 803 804 805 806 807  if bra is None: raise TypeError('bra must be an array') bra = np.asarray(bra, dtype=complex) ket = bra if ket is None else np.asarray(ket, dtype=complex) tot_norbs = _get_tot_norbs(self.syst) if bra.shape != (tot_norbs,): msg = 'vector is incorrect shape'  Joseph Weston committed Mar 10, 2017 808  msg = 'bra ' + msg if ket is None else msg  Joseph Weston committed Aug 04, 2016 809 810 811 812 813  raise ValueError(msg) elif ket.shape != (tot_norbs,): raise ValueError('ket vector is incorrect shape') result = np.zeros((self.where.shape[0],), dtype=complex)  Joseph Weston committed Feb 21, 2017 814 815  self._operate(out_data=result, bra=bra, ket=ket, args=args, params=params, op=MAT_ELS)  Joseph Weston committed Jan 12, 2017 816 817  # if everything is Hermitian then result is real if bra == ket if self.check_hermiticity and bra is ket:  Joseph Weston committed Aug 04, 2016 818  result = result.real  Joseph Weston committed Dec 15, 2016 819  return np.sum(result) if self.sum else result  Joseph Weston committed Aug 04, 2016 820 821  @cython.embedsignature  Joseph Weston committed Feb 21, 2017 822  def act(self, ket, args=(), *, params=None):  Joseph Weston committed Aug 04, 2016 823 824  """Act with the operator on a wavefunction.  Joseph Weston committed Dec 15, 2016 825 826 827  For an operator :math:Q_{iαβ} and ket :math:ψ_β this computes :math:∑_{iβ} Q_{iαβ} ψ_β.  Joseph Weston committed Aug 04, 2016 828 829  Parameters ----------  Joseph Weston committed Dec 15, 2016 830  ket : sequence of complex  Joseph Weston committed Aug 04, 2016 831 832 833  Wavefunctions defined over all the orbitals of the system. args : tuple The extra arguments to the Hamiltonian value functions and  Joseph Weston committed Feb 15, 2019 834 835  the operator onsite function. Deprecated in favor of 'params' (and mutually exclusive with it).  Joseph Weston committed Feb 21, 2017 836 837 838  params : dict, optional Dictionary of parameter names and their values. Mutually exclusive with 'args'.  Joseph Weston committed Aug 04, 2016 839 840 841  Returns -------  Joseph Weston committed Dec 15, 2016 842  Array of complex.  Joseph Weston committed Aug 04, 2016 843  """  Joseph Weston committed Feb 21, 2017 844 845 846 847  if (self._bound_onsite or self._bound_hamiltonian) and (args or params): raise ValueError("Extra arguments are already bound to this " "operator. You should call this operator " "providing neither 'args' nor 'params'.")  Joseph Weston committed Feb 15, 2019 848 849 850 851 852  if args: # deprecate_args does not play nicely with methods of cdef classes, # when used as a decorator, so we manually raise the # deprecation warning here. deprecate_args()  Joseph Weston committed Feb 21, 2017 853 854  if args and params: raise TypeError("'args' and 'params' are mutually exclusive.")  Joseph Weston committed Aug 04, 2016 855 856 857 858 859 860 861  if ket is None: raise TypeError('ket must be an array') ket = np.asarray(ket, dtype=complex) tot_norbs = _get_tot_norbs(self.syst) if ket.shape != (tot_norbs,): raise ValueError('ket vector is incorrect shape')  Viacheslav Ostroukh committed Mar 02, 2021 862  result = np.zeros((tot_norbs,), dtype=complex)  Joseph Weston committed Feb 21, 2017 863 864  self._operate(out_data=result, bra=None, ket=ket, args=args, params=params, op=ACT)  Joseph Weston committed Aug 04, 2016 865 866 867  return result @cython.embedsignature  Joseph Weston committed Feb 21, 2017 868  def bind(self, args=(), *, params=None):  Joseph Weston committed Aug 04, 2016 869 870 871 872  """Bind the given arguments to this operator. Returns a copy of this operator that does not need to be passed extra arguments when subsequently called or when using the act method.  Joseph Weston committed Feb 15, 2019 873 874 875  Providing positional arguments via 'args' is deprecated, instead provide named parameters as a dictionary via 'params'.  Joseph Weston committed Aug 04, 2016 876  """  Joseph Weston committed Feb 15, 2019 877 878 879 880 881  if args: # deprecate_args does not play nicely with methods of cdef classes, # when used as a decorator, so we manually raise the # deprecation warning here. deprecate_args()  Joseph Weston committed Feb 21, 2017 882 883  if args and params: raise TypeError("'args' and 'params' are mutually exclusive.")  Joseph Weston committed Aug 04, 2016 884 885 886 887 888  # generic creation of new instance cls = self.__class__ q = cls.__new__(cls) q.syst = self.syst q.onsite = self.onsite  Christoph Groth committed Sep 20, 2018 889  q._onsite_param_names = self._onsite_param_names  Joseph Weston committed Aug 04, 2016 890  q.where = self.where  Joseph Weston committed Jan 17, 2017 891  q.sum = self.sum  Joseph Weston committed Aug 04, 2016 892 893 894  q._site_ranges = self._site_ranges q.check_hermiticity = self.check_hermiticity if callable(self.onsite):  Joseph Weston committed Feb 21, 2017 895  q._bound_onsite = self._eval_onsites(args, params)  Joseph Weston committed Aug 04, 2016 896 897 898 899  # NOTE: subclasses should populate bound_hamiltonian if needed return q def _operate(self, complex[:] out_data, complex[:] bra, complex[:] ket,  Joseph Weston committed Feb 21, 2017 900  args, operation op, *, params=None):  Joseph Weston committed Aug 04, 2016 901 902 903 904 905 906 907 908 909 910 911 912 913  """Do an operation with the operator. Parameters ---------- out_data : ndarray Output array, zero on entry. On exit should contain the required data. What this means depends on the value of op, as does the length of the array. bra, ket : ndarray Wavefunctions defined over all the orbitals of the system. If op is ACT then bra is None. args : tuple The extra arguments to the Hamiltonian value functions and  Joseph Weston committed Feb 15, 2019 914 915  the operator onsite function. Deprecated in favor of 'params' (and mutually exclusive with it).  Joseph Weston committed Aug 04, 2016 916 917 918 919  op : operation The operation to perform. MAT_ELS: calculate matrix elements between bra and ket ACT: act on ket with the operator  Joseph Weston committed Feb 21, 2017 920 921 922  params : dict, optional Dictionary of parameter names and their values. Mutually exclusive with 'args'.  Joseph Weston committed Aug 04, 2016 923 924 925  """ raise NotImplementedError()  Joseph Weston committed Feb 21, 2017 926  cdef BlockSparseMatrix _eval_onsites(self, args, params):  Joseph Weston committed Aug 04, 2016 927 928  """Evaluate the onsite matrices on all elements of where""" assert callable(self.onsite)  Joseph Weston committed Feb 21, 2017 929  assert not (args and params)  Joseph Weston committed Aug 04, 2016 930  check_hermiticity = self.check_hermiticity  Joseph Weston committed Dec 10, 2019 931 932 933  syst = self.syst _is_vectorized = is_vectorized(syst)  Joseph Weston committed Aug 04, 2016 934   Christoph Groth committed Sep 20, 2018 935 936 937 938 939 940 941 942 943  if params: try: args = tuple(params[pn] for pn in self._onsite_param_names) except KeyError: missing = [p for p in self._onsite_param_names if p not in params] msg = ('Operator is missing required arguments: ', ', '.join(map('"{}"'.format, missing))) raise TypeError(''.join(msg))  Joseph Weston committed Feb 21, 2017 944   Joseph Weston committed Dec 10, 2019 945 946 947 948 949 950 951 952 953 954 955 956 957 958  # Evaluate many onsites at once. See _LocalOperator.__init__ # for an explanation the parameters. def eval_onsite(term_id, term_selector, which): if _is_vectorized: if term_id >= 0: (sr, _), _ = syst.subgraphs[syst.terms[term_id].subgraph] else: (_, sr), _ = syst.subgraphs[syst.terms[-term_id - 1].subgraph] else: sr, _ = term_id start_site, norbs, _ = self.syst.site_ranges[sr] # All sites selected by 'which' are part of the same site family. site_offsets = _select(self.where, which)[:, 0] - start_site data = self.onsite(sr, site_offsets, *args)  Joseph Weston committed Dec 10, 2019 959  _check_onsites(data, norbs, self.check_hermiticity)  Joseph Weston committed Dec 10, 2019 960 961 962  return data matrix_elements = _make_matrix_elements(eval_onsite, self._terms)  Joseph Weston committed Aug 04, 2016 963  offsets, norbs = _get_all_orbs(self.where, self._site_ranges)  Joseph Weston committed Dec 10, 2019 964 965  return BlockSparseMatrix(self.where, offsets, norbs, matrix_elements)  Joseph Weston committed Aug 04, 2016 966   Joseph Weston committed Feb 21, 2017 967  cdef BlockSparseMatrix _eval_hamiltonian(self, args, params):  Joseph Weston committed Aug 04, 2016 968  """Evaluate the Hamiltonian on all elements of where."""  Joseph Weston committed Dec 10, 2019 969 970 971 972  where = self.where syst = self.syst is_onsite = self.where.shape[1] == 1  Joseph Weston committed Aug 04, 2016 973 974  check_hermiticity = self.check_hermiticity  Joseph Weston committed Dec 10, 2019 975  if is_vectorized(self.syst):  Joseph Weston committed Aug 04, 2016 976   Joseph Weston committed Dec 10, 2019 977 978 979 980 981 982 983 984 985 986 987  # Evaluate many Hamiltonian elements at once. # See _LocalOperator.__init__ for an explanation the parameters. def eval_hamiltonian(term_id, term_selector, which): herm_conj = term_id < 0 assert not is_onsite or (is_onsite and not herm_conj) # onsite terms are never hermitian conjugate if herm_conj: term_id = -term_id - 1 data = syst.hamiltonian_term(term_id, term_selector, args=args, params=params) if herm_conj: data = data.conjugate().transpose(0, 2, 1)  Joseph Weston committed Dec 10, 2019 988 989 990 991 992 993 994  # Checks for data consistency (to_sr, from_sr), _ = syst.subgraphs[syst.terms[term_id].subgraph] to_norbs = syst.site_ranges[to_sr][1] from_norbs = syst.site_ranges[from_sr][1] if herm_conj: to_norbs, from_norbs = from_norbs, to_norbs _check_hams(data, to_norbs, from_norbs, is_onsite and check_hermiticity)  Joseph Weston committed Dec 10, 2019 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012  return data else: # Evaluate many Hamiltonian elements at once. # See _LocalOperator.__init__ for an explanation the parameters. def eval_hamiltonian(term_id, term_selector, which): if is_onsite: data = [ syst.hamiltonian(where[i, 0], where[i, 0], *args, params=params) for i in which ] else: data = [ syst.hamiltonian(where[i, 0], where[i, 1], *args, params=params) for i in which ]  Joseph Weston committed Dec 12, 2019 1013   Joseph Weston committed Dec 10, 2019 1014  (to_sr, from_sr) = term_id  Joseph Weston committed Dec 12, 2019 1015 1016 1017 1018 1019  to_norbs, from_norbs = syst.site_ranges[to_sr][1], syst.site_ranges[from_sr][1] expected_shape = (len(which), to_norbs, from_norbs) data = _normalize_matrix_blocks(data, expected_shape) # Checks for data consistency  Joseph Weston committed Dec 10, 2019 1020  _check_hams(data, to_norbs, from_norbs, is_onsite and check_hermiticity)  Joseph Weston committed Dec 10, 2019 1021 1022 1023 1024 1025 1026  return data matrix_elements = _make_matrix_elements(eval_hamiltonian, self._terms) offsets, norbs = _get_all_orbs(where, self._site_ranges) return BlockSparseMatrix(where, offsets, norbs, matrix_elements)  Joseph Weston committed Aug 04, 2016 1027   Joseph Weston committed Mar 17, 2017 1028 1029 1030  def __getstate__(self): return ( (self.check_hermiticity, self.sum),  Christoph Groth committed Sep 20, 2018 1031  (self.syst, self.onsite, self._onsite_param_names),  Joseph Weston committed Mar 17, 2017 1032  tuple(map(np.asarray, (self.where, self._site_ranges))),  Joseph Weston committed Dec 10, 2019 1033  (self._terms,),  Joseph Weston committed Mar 17, 2017 1034 1035 1036 1037 1038  (self._bound_onsite, self._bound_hamiltonian), ) def __setstate__(self, state): ((self.check_hermiticity, self.sum),  Christoph Groth committed Sep 20, 2018 1039  (self.syst, self.onsite, self._onsite_param_names),  Joseph Weston committed Mar 17, 2017 1040  (self.where, self._site_ranges),  Joseph Weston committed Dec 10, 2019 1041  (self._terms,),  Joseph Weston committed Mar 17, 2017 1042 1043 1044  (self._bound_onsite, self._bound_hamiltonian), ) = state  Joseph Weston committed Aug 04, 2016 1045 1046 1047 1048  cdef class Density(_LocalOperator): """An operator for calculating general densities.  Joseph Weston committed Feb 02, 2017 1049  An instance of this class can be called like a function to evaluate the  Joseph Weston committed Dec 08, 2017 1050 1051  expectation value with a wavefunction. See ~kwant.operator.Density.__call__ for details.  Joseph Weston committed Aug 04, 2016 1052 1053 1054 1055 1056 1057 1058 1059 1060  Parameters ---------- syst : ~kwant.system.System onsite : scalar or square matrix or dict or callable The onsite matrix that defines the operator. If a dict is given, it maps from site families to square matrices. If a function is given it must take the same arguments as the onsite Hamiltonian functions of the system.  Joseph Weston committed Nov 13, 2019 1061  where : sequence of int or ~kwant.system.Site, or callable, optional  Joseph Weston committed Aug 04, 2016 1062 1063  Where to evaluate the operator. If syst is not a finalized Builder, then this should be a sequence of integers. If a function is provided,  Joseph Weston committed Nov 13, 2019 1064  it should take a single int or ~kwant.system.Site (if syst is  Joseph Weston committed Aug 04, 2016 1065 1066 1067 1068 1069 1070  a finalized builder) and return True or False. If not provided, the operator will be calculated over all sites in the system. check_hermiticity: bool Check whether the provided onsite is Hermitian. If it is not Hermitian, then an error will be raised when the operator is evaluated.  Joseph Weston committed Feb 02, 2017 1071  sum : bool, default: False  Joseph Weston committed Dec 08, 2017 1072 1073 1074  If True, then calling this operator will return a single scalar, otherwise a vector will be returned (see ~kwant.operator.Density.__call__ for details).  Joseph Weston committed Aug 04, 2016 1075   Joseph Weston committed Feb 02, 2017 1076 1077 1078 1079 1080 1081 1082  Notes ----- In general, if there is a certain "density" (e.g. charge or spin) that is represented by a square matrix :math:M_i associated with each site :math:i then an instance of this class represents the tensor :math:Q_{iαβ} which is equal to :math:M_i when α and β are orbitals on site :math:i, and zero otherwise.  Joseph Weston committed Aug 04, 2016 1083 1084 1085  """ @cython.embedsignature  Joseph Weston committed Dec 15, 2016 1086 1087 1088 1089  def __init__(self, syst, onsite=1, where=None, *, check_hermiticity=True, sum=False): where = _normalize_site_where(syst, where) super().__init__(syst, onsite, where,  Joseph Weston committed Jan 17, 2017 1090  check_hermiticity=check_hermiticity, sum=sum)  Joseph Weston committed Aug 04, 2016 1091 1092 1093 1094  @cython.boundscheck(False) @cython.wraparound(False) def _operate(self, complex[:] out_data, complex[:] bra, complex[:] ket,  Joseph Weston committed Feb 21, 2017 1095  args, operation op, *, params=None):  Joseph Weston committed Aug 04, 2016 1096 1097 1098 1099  matrix = ta.matrix cdef int unique_onsite = not callable(self.onsite) # prepare onsite matrices cdef complex[:, :] _tmp_mat  Joseph Weston committed Jan 12, 2017 1100  cdef complex *M_a = NULL  Joseph Weston committed Aug 04, 2016 1101  cdef BlockSparseMatrix M_a_blocks  Joseph Weston committed Jan 12, 2017 1102   Joseph Weston committed Aug 04, 2016 1103 1104 1105  if unique_onsite: _tmp_mat = self.onsite M_a = &_tmp_mat[0, 0]  Joseph Weston committed Jan 12, 2017 1106 1107  elif self._bound_onsite: M_a_blocks = self._bound_onsite  Joseph Weston committed Aug 04, 2016 1108  else:  Joseph Weston committed Feb 21, 2017 1109  M_a_blocks = self._eval_onsites(args, params)  Joseph Weston committed Jan 12, 2017 1110   Joseph Weston committed Aug 04, 2016 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137  # loop-local variables cdef gint a, a_s, a_norbs cdef gint i, j, w cdef complex tmp, bra_conj ### loop over sites for w in range(self.where.shape[0]): ### get the next site, start orbital and number of orbitals a = self.where[w, 0] _get_orbs(self._site_ranges, a, &a_s, &a_norbs) ### get the next onsite matrix, if necessary if not unique_onsite: M_a = M_a_blocks.get(w) ### do the actual calculation if op == MAT_ELS: tmp = 0 for i in range(a_norbs): for j in range(a_norbs): tmp += (bra[a_s + i].conjugate() * M_a[i * a_norbs + j] * ket[a_s + j]) out_data[w] = tmp elif op == ACT: for i in range(a_norbs): tmp = 0 for j in range(a_norbs): tmp += M_a[i * a_norbs + j] * ket[a_s + j] out_data[a_s + i] = out_data[a_s + i] + tmp  Anton Akhmerov committed Jan 22, 2017 1138 1139 1140  @cython.boundscheck(False) @cython.wraparound(False) @cython.cdivision(True)  Joseph Weston committed Feb 02, 2017 1141  @cython.embedsignature  Joseph Weston committed Feb 21, 2017 1142  def tocoo(self, args=(), *, params=None):  Joseph Weston committed Feb 15, 2019 1143 1144 1145 1146 1147  """Convert the operator to coordinate format sparse matrix. Providing positional arguments via 'args' is deprecated, instead provide named parameters as a dictionary via 'params'. """  Anton Akhmerov committed Jan 22, 2017 1148 1149 1150  cdef int blk, blk_size, n_blocks, n, k = 0 cdef int [:, :] offsets, shapes cdef int [:] row, col  Joseph Weston committed Feb 21, 2017 1151  if self._bound_onsite and (args or params):  Anton Akhmerov committed Jan 22, 2017 1152 1153  raise ValueError("Extra arguments are already bound to this " "operator. You should call this operator "  Joseph Weston committed Feb 21, 2017 1154  "providing neither 'args' nor 'params'.")  Joseph Weston committed Feb 15, 2019 1155 1156 1157 1158 1159  if args: # deprecate_args does not play nicely with methods of cdef classes, # when used as a decorator, so we manually raise the # deprecation warning here. deprecate_args()  Joseph Weston committed Feb 21, 2017 1160 1161  if args and params: raise TypeError("'args' and 'params' are mutually exclusive.")  Anton Akhmerov committed Jan 22, 2017 1162 1163 1164 1165 1166 1167 1168 1169 1170  if not callable(self.onsite): offsets = _get_all_orbs(self.where, self._site_ranges)[0] n_blocks = len(self.where) shapes = np.asarray(np.resize([self.onsite.shape], (n_blocks, 2)), gint_dtype) data = np.asarray(self.onsite).flatten() data = np.resize(data, [len(data) * n_blocks]) else:  Jörg Behrmann committed Feb 02, 2017 1171 1172 1173  if self._bound_onsite is not None: onsite_matrix = self._bound_onsite else:  Joseph Weston committed Feb 21, 2017 1174  onsite_matrix = self._eval_onsites(args, params)  Anton Akhmerov committed Jan 22, 2017 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192  data = onsite_matrix.data offsets = np.asarray(onsite_matrix.block_offsets) shapes = np.asarray(onsite_matrix.block_shapes) row = np.empty(len(data), gint_dtype) col = np.empty(len(data), gint_dtype) for blk in range(len(offsets)): blk_size = shapes[blk, 0] * shapes[blk, 1] for n in range(blk_size): row[k] = offsets[blk, 0] + n // shapes[blk, 1] col[k] = offsets[blk, 1] + n % shapes[blk, 1] k += 1 norbs = _get_tot_norbs(self.syst) return coo_matrix((np.asarray(data), (np.asarray(row), np.asarray(col))), shape=(norbs, norbs))  Joseph Weston committed Aug 04, 2016 1193 1194  cdef class Current(_LocalOperator):  Joseph Weston committed May 12, 2017 1195  r"""An operator for calculating general currents.  Joseph Weston committed Aug 04, 2016 1196   Joseph Weston committed Feb 02, 2017 1197  An instance of this class can be called like a function to evaluate the  Joseph Weston committed Dec 08, 2017 1198 1199  expectation value with a wavefunction. See ~kwant.operator.Current.__call__ for details.  Joseph Weston committed Aug 04, 2016 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209  Parameters ---------- syst : ~kwant.system.System onsite : scalar or square matrix or dict or callable The onsite matrix that defines the density from which this current is derived. If a dict is given, it maps from site families to square matrices (scalars are allowed if the site family has 1 orbital per site). If a function is given it must take the same arguments as the onsite Hamiltonian functions of the system.  Joseph Weston committed Nov 13, 2019 1210  where : sequence of pairs of int or ~kwant.system.Site, or callable, optional  Joseph Weston committed Aug 04, 2016 1211 1212 1213  Where to evaluate the operator. If syst is not a finalized Builder, then this should be a sequence of pairs of integers. If a function is provided, it should take a pair of integers or a pair of  Joseph Weston committed Nov 13, 2019 1214  ~kwant.system.Site (if syst is a finalized builder) and return  Joseph Weston committed Aug 04, 2016 1215 1216 1217 1218 1219 1220  True or False. If not provided, the operator will be calculated over all hoppings in the system. check_hermiticity : bool Check whether the provided onsite is Hermitian. If it is not Hermitian, then an error will be raised when the operator is evaluated.  Joseph Weston committed Feb 02, 2017 1221  sum : bool, default: False  Joseph Weston committed Dec 08, 2017 1222 1223 1224  If True, then calling this operator will return a single scalar, otherwise a vector will be returned (see ~kwant.operator.Current.__call__ for details).  Joseph Weston committed Aug 04, 2016 1225   Joseph Weston committed Feb 02, 2017 1226 1227 1228 1229 1230 1231  Notes ----- In general, if there is a certain "density" (e.g. charge or spin) that is represented by a square matrix :math:M_i associated with each site :math:i and :math:H_{ij} is the hopping Hamiltonian from site :math:j to site i, then an instance of this class represents the tensor  Joseph Weston committed May 12, 2017 1232 1233 1234  :math:J_{ijαβ} which is equal to :math:i\left[(H_{ij})^† M_i - M_i H_{ij}\right] when α and β are orbitals on sites :math:i and :math:j respectively, and zero otherwise.  Joseph Weston committed Feb 02, 2017 1235 1236 1237  The tensor :math:J_{ijαβ} will also be referred to as :math:Q_{nαβ}, where :math:n is the index of hopping :math:(i, j) in where.  Joseph Weston committed Aug 04, 2016 1238 1239 1240  """ @cython.embedsignature  Joseph Weston committed Dec 15, 2016 1241 1242 1243 1244  def __init__(self, syst, onsite=1, where=None, *, check_hermiticity=True, sum=False): where = _normalize_hopping_where(syst, where) super().__init__(syst, onsite, where,  Joseph Weston committed Jan 17, 2017 1245  check_hermiticity=check_hermiticity, sum=sum)  Joseph Weston committed Aug 04, 2016 1246 1247  @cython.embedsignature  Joseph Weston committed Feb 21, 2017 1248  def bind(self, args=(), *, params=None):  Joseph Weston committed Aug 04, 2016 1249 1250 1251 1252  """Bind the given arguments to this operator. Returns a copy of this operator that does not need to be passed extra arguments when subsequently called or when using the act method.  Joseph Weston committed Feb 15, 2019 1253 1254 1255  Providing positional arguments via 'args' is deprecated, instead provide named parameters as a dictionary via 'params'.  Joseph Weston committed Aug 04, 2016 1256  """  Joseph Weston committed Feb 21, 2017 1257 1258  q = super().bind(args, params=params) q._bound_hamiltonian = self._eval_hamiltonian(args, params)  Joseph Weston committed Aug 04, 2016 1259 1260 1261 1262 1263  return q @cython.boundscheck(False) @cython.wraparound(False) def _operate(self, complex[:] out_data, complex[:] bra, complex[:] ket,  Joseph Weston committed Feb 21, 2017 1264  args, operation op, *, params=None):  Joseph Weston committed Aug 04, 2016 1265 1266 1267  # prepare onsite matrices and hamiltonians cdef int unique_onsite = not callable(self.onsite) cdef complex[:, :] _tmp_mat  Joseph Weston committed Jan 12, 2017 1268 1269  cdef complex *M_a = NULL cdef complex *H_ab = NULL  Joseph Weston committed Aug 04, 2016 1270  cdef BlockSparseMatrix M_a_blocks, H_ab_blocks  Joseph Weston committed Jan 12, 2017 1271   Joseph Weston committed Aug 04, 2016 1272 1273 1274  if unique_onsite: _tmp_mat = self.onsite M_a = &_tmp_mat[0, 0]  Joseph Weston committed Jan 12, 2017 1275 1276 1277  elif self._bound_onsite: M_a_blocks = self._bound_onsite else:  Joseph Weston committed Feb 21, 2017 1278  M_a_blocks = self._eval_onsites(args, params)  Joseph Weston committed Jan 12, 2017 1279 1280 1281  if self._bound_hamiltonian: H_ab_blocks = self._bound_hamiltonian  Joseph Weston committed Aug 04, 2016 1282  else:  Joseph Weston committed Feb 21, 2017 1283  H_ab_blocks = self._eval_hamiltonian(args, params)  Joseph Weston committed Jan 12, 2017 1284   Joseph Weston committed Aug 04, 2016 1285 1286 1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306 1307 1308 1309 1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321 1322 1323 1324 1325 1326 1327 1328  # main loop cdef gint a, a_s, a_norbs, b, b_s, b_norbs cdef gint i, j, k, w cdef complex tmp for w in range(self.where.shape[0]): ### get the next hopping's start orbitals and numbers of orbitals a_s = H_ab_blocks.block_offsets[w, 0] b_s = H_ab_blocks.block_offsets[w, 1] a_norbs = H_ab_blocks.block_shapes[w, 0] b_norbs = H_ab_blocks.block_shapes[w, 1] ### get the next onsite and Hamiltonian matrices H_ab = H_ab_blocks.get(w) if not unique_onsite: M_a = M_a_blocks.get(w) ### do the actual calculation if op == MAT_ELS: tmp = 0 for i in range(b_norbs): for j in range(a_norbs): for k in range(a_norbs): tmp += (bra[b_s + i].conjugate() * H_ab[j * b_norbs + i].conjugate() * M_a[j * a_norbs + k] * ket[a_s + k] - bra[a_s + j].conjugate() * M_a[j * a_norbs + k] * H_ab[k * b_norbs + i] * ket[b_s + i]) out_data[w] = 1j * tmp elif op == ACT: for i in range(b_norbs): for j in range(a_norbs): for k in range(a_norbs): out_data[b_s + i] = ( out_data[b_s + i] + 1j * H_ab[j * b_norbs + i].conjugate() * M_a[j * a_norbs + k] * ket[a_s + k]) out_data[a_s + j] = ( out_data[a_s + j] - 1j * M_a[j * a_norbs + k] * H_ab[k * b_norbs + i] * ket[b_s + i]) cdef class Source(_LocalOperator): """An operator for calculating general sources.  Joseph Weston committed Feb 02, 2017 1329  An instance of this class can be called like a function to evaluate the  Joseph Weston committed Dec 08, 2017 1330 1331  expectation value with a wavefunction. See ~kwant.operator.Source.__call__ for details.  Joseph Weston committed Aug 04, 2016 1332 1333 1334 1335 1336 1337 1338 1339 1340 1341  Parameters ---------- syst : ~kwant.system.System onsite : scalar or square matrix or dict or callable The onsite matrix that defines the density from which this source is defined. If a dict is given, it maps from site families to square matrices (scalars are allowed if the site family has 1 orbital per site). If a function is given it must take the same arguments as the onsite Hamiltonian functions of the system.  Joseph Weston committed Nov 13, 2019 1342  where : sequence of int or ~kwant.system.Site, or callable, optional  Joseph Weston committed Aug 04, 2016 1343 1344  Where to evaluate the operator. If syst is not a finalized Builder, then this should be a sequence of integers. If a function is provided,  Joseph Weston committed Nov 13, 2019 1345  it should take a single int or ~kwant.system.Site (if syst is  Joseph Weston committed Aug 04, 2016 1346 1347 1348 1349 1350 1351  a finalized builder) and return True or False. If not provided, the operator will be calculated over all sites in the system. check_hermiticity : bool Check whether the provided onsite is Hermitian. If it is not Hermitian, then an error will be raised when the operator is evaluated.  Joseph Weston committed Feb 02, 2017 1352  sum : bool, default: False  Joseph Weston committed Dec 08, 2017 1353