Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright 2011-2019 Kwant authors.
2#
3# This file is part of Kwant. It is subject to the license terms in the file
4# LICENSE.rst found in the top-level directory of this distribution and at
5# https://kwant-project.org/license. A list of Kwant authors can be found in
6# the file AUTHORS.rst at the top-level directory of this distribution and at
7# https://kwant-project.org/authors.
9cimport cython
10import tinyarray as ta
11import numpy as np
12from scipy import sparse as sp
13from itertools import chain
14import types
15import bisect
17from .graph.core cimport CGraph, gintArraySlice
18from .graph.defs cimport gint
19from .graph.defs import gint_dtype
20from ._common import deprecate_args
23### Non-vectorized methods
25msg = ('Hopping from site {0} to site {1} does not match the '
26 'dimensions of onsite Hamiltonians of these sites.')
28@cython.boundscheck(False)
29def make_sparse(ham, args, params, CGraph gr, diag,
30 gint [:] from_sites, n_by_to_site,
31 gint [:] to_norb, gint [:] to_off,
32 gint [:] from_norb, gint [:] from_off):
33 """For internal use by hamiltonian_submatrix."""
34 cdef gintArraySlice nbors
35 cdef gint n_fs, fs, n_ts, ts
36 cdef gint i, j, num_entries
37 cdef complex [:, :] h
38 cdef gint [:, :] rows_cols
39 cdef complex [:] data
40 cdef complex value
42 matrix = ta.matrix
44 # Calculate the data size.
45 num_entries = 0
46 for n_fs in range(len(from_sites)):
47 fs = from_sites[n_fs]
48 if fs in n_by_to_site:
49 num_entries += from_norb[n_fs] * from_norb[n_fs]
50 nbors = gr.out_neighbors(fs)
51 for ts in nbors.data[:nbors.size]:
52 if ts in n_by_to_site:
53 n_ts = n_by_to_site[ts]
54 num_entries += to_norb[n_ts] * from_norb[n_fs]
56 rows_cols = np.empty((2, num_entries), gint_dtype)
57 data = np.empty(num_entries, complex)
59 cdef gint k = 0
60 for n_fs in range(len(from_sites)):
61 fs = from_sites[n_fs]
62 if fs in n_by_to_site:
63 n_ts = n_by_to_site[fs]
64 h = diag[n_fs]
65 if not (h.shape[0] == h.shape[1] == from_norb[n_fs]):
66 raise ValueError(msg.format(fs, fs))
67 for i in range(h.shape[0]):
68 for j in range(h.shape[1]):
69 value = h[i, j]
70 if value != 0:
71 data[k] = value
72 rows_cols[0, k] = i + to_off[n_ts]
73 rows_cols[1, k] = j + from_off[n_fs]
74 k += 1
76 nbors = gr.out_neighbors(fs)
77 for ts in nbors.data[:nbors.size]:
78 if ts not in n_by_to_site:
79 continue
80 n_ts = n_by_to_site[ts]
81 h = matrix(ham(ts, fs, *args, params=params), complex)
82 if h.shape[0] != to_norb[n_ts] or h.shape[1] != from_norb[n_fs]:
83 raise ValueError(msg.format(fs, ts))
84 for i in range(h.shape[0]):
85 for j in range(h.shape[1]):
86 value = h[i, j]
87 if value != 0:
88 data[k] = value
89 rows_cols[0, k] = i + to_off[n_ts]
90 rows_cols[1, k] = j + from_off[n_fs]
91 k += 1
93 return sp.coo_matrix((data[:k], rows_cols[:, :k]),
94 shape=(to_off[-1], from_off[-1]))
97@cython.boundscheck(False)
98def make_sparse_full(ham, args, params, CGraph gr, diag,
99 gint [:] to_norb, gint [:] to_off,
100 gint [:] from_norb, gint [:] from_off):
101 """For internal use by hamiltonian_submatrix."""
102 cdef gintArraySlice nbors
103 cdef gint n, fs, ts
104 cdef gint i, j, num_entries
105 cdef complex [:, :] h
106 cdef gint [:, :] rows_cols
107 cdef complex [:] data
108 cdef complex value
110 matrix = ta.matrix
111 n = gr.num_nodes
113 # Calculate the data size.
114 num_entries = 0
115 for fs in range(n):
116 num_entries += from_norb[fs] * from_norb[fs]
117 nbors = gr.out_neighbors(fs)
118 for ts in nbors.data[:nbors.size]:
119 if fs < ts:
120 num_entries += 2 * to_norb[ts] * from_norb[fs]
122 rows_cols = np.empty((2, num_entries), gint_dtype)
123 data = np.empty(num_entries, complex)
125 cdef gint k = 0
126 for fs in range(n):
127 h = diag[fs]
128 if not (h.shape[0] == h.shape[1] == from_norb[fs]):
129 raise ValueError(msg.format(fs, fs))
130 for i in range(h.shape[0]):
131 for j in range(h.shape[1]):
132 value = h[i, j]
133 if value != 0:
134 data[k] = value
135 rows_cols[0, k] = i + to_off[fs]
136 rows_cols[1, k] = j + from_off[fs]
137 k += 1
139 nbors = gr.out_neighbors(fs)
140 for ts in nbors.data[:nbors.size]:
141 if ts < fs:
142 continue
143 h = matrix(ham(ts, fs, *args, params=params), complex)
144 if h.shape[0] != to_norb[ts] or h.shape[1] != from_norb[fs]:
145 raise ValueError(msg.format(fs, ts))
146 for i in range(h.shape[0]):
147 for j in range(h.shape[1]):
148 value = h[i, j]
149 if value != 0:
150 data[k] = value
151 data[k + 1] = h[i, j].conjugate()
152 rows_cols[1, k + 1] = rows_cols[0, k] = i + to_off[ts]
153 rows_cols[0, k + 1] = rows_cols[1, k] = j + from_off[fs]
154 k += 2
156 return sp.coo_matrix((data[:k], rows_cols[:, :k]),
157 shape=(to_off[-1], from_off[-1]))
160@cython.boundscheck(False)
161def make_dense(ham, args, params, CGraph gr, diag,
162 gint [:] from_sites, n_by_to_site,
163 gint [:] to_norb, gint [:] to_off,
164 gint [:] from_norb, gint [:] from_off):
165 """For internal use by hamiltonian_submatrix."""
166 cdef gintArraySlice nbors
167 cdef gint n_fs, fs, n_ts, ts
168 cdef complex [:, :] h_sub_view
169 cdef complex [:, :] h
171 matrix = ta.matrix
173 h_sub = np.zeros((to_off[-1], from_off[-1]), complex)
174 h_sub_view = h_sub
175 for n_fs in range(len(from_sites)):
176 fs = from_sites[n_fs]
177 if fs in n_by_to_site:
178 n_ts = n_by_to_site[fs]
179 h = diag[n_fs]
180 if not (h.shape[0] == h.shape[1] == from_norb[n_fs]):
181 raise ValueError(msg.format(fs, fs))
182 h_sub_view[to_off[n_ts] : to_off[n_ts + 1],
183 from_off[n_fs] : from_off[n_fs + 1]] = h
185 nbors = gr.out_neighbors(fs)
186 for ts in nbors.data[:nbors.size]:
187 if ts not in n_by_to_site:
188 continue
189 n_ts = n_by_to_site[ts]
190 h = matrix(ham(ts, fs, *args, params=params), complex)
191 if h.shape[0] != to_norb[n_ts] or h.shape[1] != from_norb[n_fs]:
192 raise ValueError(msg.format(fs, ts))
193 h_sub_view[to_off[n_ts] : to_off[n_ts + 1],
194 from_off[n_fs] : from_off[n_fs + 1]] = h
195 return h_sub
198@cython.boundscheck(False)
199def make_dense_full(ham, args, params, CGraph gr, diag,
200 gint [:] to_norb, gint [:] to_off,
201 gint [:] from_norb, gint [:] from_off):
202 """For internal use by hamiltonian_submatrix."""
203 cdef gintArraySlice nbors
204 cdef gint n, fs, ts
205 cdef complex [:, :] h_sub_view, h, h_herm
207 matrix = ta.matrix
208 n = gr.num_nodes
210 h_sub = np.zeros((to_off[-1], from_off[-1]), complex)
211 h_sub_view = h_sub
212 for fs in range(n):
213 h = diag[fs]
214 if not (h.shape[0] == h.shape[1] == from_norb[fs]):
215 raise ValueError(msg.format(fs, fs))
216 h_sub_view[to_off[fs] : to_off[fs + 1],
217 from_off[fs] : from_off[fs + 1]] = h
219 nbors = gr.out_neighbors(fs)
220 for ts in nbors.data[:nbors.size]:
221 if ts < fs:
222 continue
223 h = mat = matrix(ham(ts, fs, *args, params=params), complex)
224 h_herm = mat.transpose().conjugate()
225 if h.shape[0] != to_norb[ts] or h.shape[1] != from_norb[fs]:
226 raise ValueError(msg.format(fs, ts))
227 h_sub_view[to_off[ts] : to_off[ts + 1],
228 from_off[fs] : from_off[fs + 1]] = h
229 h_sub_view[from_off[fs] : from_off[fs + 1],
230 to_off[ts] : to_off[ts + 1]] = h_herm
231 return h_sub
234@deprecate_args
235@cython.binding(True)
236@cython.embedsignature(True)
237def hamiltonian_submatrix(self, args=(), to_sites=None, from_sites=None,
238 sparse=False, return_norb=False, *, params=None):
239 """Return a submatrix of the system Hamiltonian.
241 Parameters
242 ----------
243 args : tuple, defaults to empty
244 Positional arguments to pass to the ``hamiltonian`` method. Mutually
245 exclusive with 'params'.
246 to_sites : sequence of sites or None (default)
247 from_sites : sequence of sites or None (default)
248 sparse : bool
249 Whether to return a sparse or a dense matrix. Defaults to ``False``.
250 return_norb : bool
251 Whether to return arrays of numbers of orbitals. Defaults to ``False``.
252 params : dict, optional
253 Dictionary of parameter names and their values. Mutually exclusive
254 with 'args'.
256 Returns
257 -------
258 hamiltonian_part : numpy.ndarray or scipy.sparse.coo_matrix
259 Submatrix of Hamiltonian of the system.
260 to_norb : array of integers
261 Numbers of orbitals on each site in to_sites. Only returned when
262 ``return_norb`` is true.
263 from_norb : array of integers
264 Numbers of orbitals on each site in from_sites. Only returned when
265 ``return_norb`` is true.
267 Notes
268 -----
269 The returned submatrix contains all the Hamiltonian matrix elements
270 from ``from_sites`` to ``to_sites``. The default for ``from_sites`` and
271 ``to_sites`` is ``None`` which means to use all sites of the system in the
272 order in which they appear.
273 """
274 cdef gint [:] to_norb, from_norb
275 cdef gint site, n_site, n
277 ham = self.hamiltonian
278 n = self.graph.num_nodes
279 matrix = ta.matrix
281 if from_sites is None:
282 diag = n * [None]
283 from_norb = np.empty(n, gint_dtype)
284 for site in range(n):
285 diag[site] = h = matrix(ham(site, site, *args, params=params),
286 complex)
287 from_norb[site] = h.shape[0]
288 else:
289 diag = len(from_sites) * [None]
290 from_norb = np.empty(len(from_sites), gint_dtype)
291 for n_site, site in enumerate(from_sites):
292 if site < 0 or site >= n:
293 raise IndexError('Site number out of range.')
294 diag[n_site] = h = matrix(ham(site, site, *args, params=params),
295 complex)
296 from_norb[n_site] = h.shape[0]
297 from_off = np.empty(from_norb.shape[0] + 1, gint_dtype)
298 from_off[0] = 0
299 from_off[1 :] = np.cumsum(from_norb)
301 if to_sites is from_sites:
302 to_norb = from_norb
303 to_off = from_off
304 else:
305 if to_sites is None:
306 to_norb = np.empty(n, gint_dtype)
307 for site in range(n):
308 h = matrix(ham(site, site, *args, params=params), complex)
309 to_norb[site] = h.shape[0]
310 else:
311 to_norb = np.empty(len(to_sites), gint_dtype)
312 for n_site, site in enumerate(to_sites):
313 if site < 0 or site >= n:
314 raise IndexError('Site number out of range.')
315 h = matrix(ham(site, site, *args, params=params), complex)
316 to_norb[n_site] = h.shape[0]
317 to_off = np.empty(to_norb.shape[0] + 1, gint_dtype)
318 to_off[0] = 0
319 to_off[1 :] = np.cumsum(to_norb)
322 if to_sites is from_sites is None:
323 func = make_sparse_full if sparse else make_dense_full
324 mat = func(ham, args, params, self.graph, diag, to_norb, to_off,
325 from_norb, from_off)
326 else:
327 if to_sites is None:
328 to_sites = np.arange(n, dtype=gint_dtype)
329 n_by_to_site = dict((site, site) for site in to_sites)
330 else:
331 n_by_to_site = dict((site, n_site)
332 for n_site, site in enumerate(to_sites))
334 if from_sites is None:
335 from_sites = np.arange(n, dtype=gint_dtype)
336 else:
337 from_sites = np.asarray(from_sites, gint_dtype)
339 func = make_sparse if sparse else make_dense
340 mat = func(ham, args, params, self.graph, diag, from_sites,
341 n_by_to_site, to_norb, to_off, from_norb, from_off)
342 return (mat, to_norb, from_norb) if return_norb else mat
345### Vectorized methods
348_shape_error_msg = (
349 "The following hoppings have matrix elements of incompatible shape "
350 "with other matrix elements: {}"
351)
354@cython.boundscheck(False)
355def _vectorized_make_sparse(subgraphs, hams, long [:] norbs, long [:] orb_offsets,
356 long [:] site_offsets, shape=None):
357 ndata = sum(h.shape[0] * h.shape[1] * h.shape[2] for h in hams)
359 cdef long [:] rows_view, cols_view
360 cdef complex [:] data_view
362 rows = np.empty((ndata,), dtype=long)
363 cols = np.empty((ndata,), dtype=long)
364 data = np.empty((ndata,), dtype=complex)
365 rows_view = rows
366 cols_view = cols
367 data_view = data
369 cdef long i, j, k, m, N, M, P, to_off, from_off,\
370 ta, fa, to_norbs, from_norbs
371 cdef const long [:] ts_offs, fs_offs
372 cdef complex [:, :, :] h
374 m = 0
375 # This outer loop zip() is pure Python, but that's ok, as it
376 # has very few entries and the inner loops are fully vectorized
377 for ((ta, fa), (ts_offs, fs_offs)), h in zip(subgraphs, hams):
378 N = h.shape[0]
379 M = h.shape[1]
380 P = h.shape[2]
382 if norbs[ta] != M or norbs[fa] != P:
383 to_sites = site_offsets[ta] + np.array(ts_offs)
384 from_sites = site_offsets[fa] + np.array(fs_offs)
385 hops = np.array([to_sites, from_sites]).transpose()
386 raise ValueError(_shape_error_msg.format(hops))
388 for i in range(N):
389 to_off = orb_offsets[ta] + norbs[ta] * ts_offs[i]
390 from_off = orb_offsets[fa] + norbs[fa] * fs_offs[i]
391 for j in range(M):
392 for k in range(P):
393 rows_view[m] = to_off + j
394 cols_view[m] = from_off + k
395 data_view[m] = h[i, j, k]
396 m += 1
398 if shape is None:
399 shape = (orb_offsets[-1], orb_offsets[-1])
401 return sp.coo_matrix((data, (rows, cols)), shape=shape)
404@cython.boundscheck(False)
405def _vectorized_make_dense(subgraphs, hams, long [:] norbs, long [:] orb_offsets,
406 long [:] site_offsets, shape=None):
407 if shape is None:
408 shape = (orb_offsets[-1], orb_offsets[-1])
409 mat = np.zeros(shape, dtype=complex)
410 cdef complex [:, :] mat_view
411 mat_view = mat
413 cdef long i, j, k, N, M, P, to_off, from_off,\
414 ta, fa, to_norbs, from_norbs
415 cdef const long [:] ts_offs, fs_offs
416 cdef complex [:, :, :] h
418 # This outer loop zip() is pure Python, but that's ok, as it
419 # has very few entries and the inner loops are fully vectorized
420 for ((ta, fa), (ts_offs, fs_offs)), h in zip(subgraphs, hams):
421 N = h.shape[0]
422 M = h.shape[1]
423 P = h.shape[2]
425 if norbs[ta] != M or norbs[fa] != P:
426 to_sites = site_offsets[ta] + np.array(ts_offs)
427 from_sites = site_offsets[fa] + np.array(fs_offs)
428 hops = np.array([to_sites, from_sites]).transpose()
429 raise ValueError(_shape_error_msg.format(hops))
431 for i in range(N):
432 to_off = orb_offsets[ta] + norbs[ta] * ts_offs[i]
433 from_off = orb_offsets[fa] + norbs[fa] * fs_offs[i]
434 for j in range(M):
435 for k in range(P):
436 mat_view[to_off + j, from_off + k] = h[i, j, k]
437 return mat
440def _expand_norbs(compressed_norbs, site_offsets):
441 "Return norbs per site, given norbs per site array."
442 norbs = np.empty(site_offsets[-1], int)
443 for start, stop, norb in zip(site_offsets, site_offsets[1:],
444 compressed_norbs):
445 norbs[start:stop] = norb
446 return norbs
449def _reverse_subgraph(subgraph):
450 (to_sa, from_sa), (to_off, from_off) = subgraph
451 return ((from_sa, to_sa), (from_off, to_off))
454@deprecate_args
455@cython.binding(True)
456def vectorized_hamiltonian_submatrix(self, args=(), sparse=False,
457 return_norb=False, *, params=None):
458 """Return The system Hamiltonian.
460 Parameters
461 ----------
462 args : tuple, defaults to empty
463 Positional arguments to pass to ``hamiltonian_term``. Mutually
464 exclusive with 'params'.
465 sparse : bool
466 Whether to return a sparse or a dense matrix. Defaults to ``False``.
467 return_norb : bool
468 Whether to return arrays of numbers of orbitals. Defaults to ``False``.
469 params : dict, optional
470 Dictionary of parameter names and their values. Mutually exclusive
471 with 'args'.
473 Returns
474 -------
475 hamiltonian_part : numpy.ndarray or scipy.sparse.coo_matrix
476 The Hamiltonian of the system.
477 norb : array of integers
478 Numbers of orbitals on each site. Only returned when ``return_norb``
479 is true.
481 Notes
482 -----
483 Providing positional arguments via 'args' is deprecated,
484 instead, provide named parameters as a dictionary via 'params'.
485 """
487 site_offsets = np.cumsum([0] + [len(arr) for arr in self.site_arrays])
489 subgraphs = [
490 self.subgraphs[t.subgraph]
491 for t in self.terms
492 ]
493 # Add Hermitian conjugates
494 subgraphs += [
495 _reverse_subgraph(self.subgraphs[t.subgraph])
496 for t in self.terms
497 if t.hermitian
498 ]
500 hams = [
501 self.hamiltonian_term(n, args=args, params=params)
502 for n, _ in enumerate(self.terms)
503 ]
504 # Add Hermitian conjugates
505 hams += [
506 ham.conjugate().transpose(0, 2, 1)
507 for ham, t in zip(hams, self.terms)
508 if t.hermitian
509 ]
511 _, norbs, orb_offsets = self.site_ranges.transpose()
513 func = _vectorized_make_sparse if sparse else _vectorized_make_dense
514 mat = func(subgraphs, hams, norbs, orb_offsets, site_offsets)
516 if return_norb:
517 return (mat, _expand_norbs(norbs, site_offsets))
518 else:
519 return mat
522@deprecate_args
523@cython.binding(True)
524def vectorized_cell_hamiltonian(self, args=(), sparse=False, *, params=None):
525 """Hamiltonian of a single cell of the infinite system.
527 Providing positional arguments via 'args' is deprecated,
528 instead, provide named parameters as a dictionary via 'params'.
529 """
531 site_offsets = np.cumsum([0] + [len(arr) for arr in self.site_arrays])
532 # Site array where next cell starts
533 next_cell = bisect.bisect(site_offsets, self.cell_size) - 1
535 def inside_fd(term):
536 return all(s == 0 for s in term.symmetry_element)
538 cell_terms = [
539 n for n, t in enumerate(self.terms)
540 if inside_fd(t)
541 ]
543 subgraphs = [
544 self.subgraphs[self.terms[n].subgraph]
545 for n in cell_terms
546 ]
547 # Add Hermitian conjugates
548 subgraphs += [
549 _reverse_subgraph(self.subgraphs[self.terms[n].subgraph])
550 for n in cell_terms
551 if self.terms[n].hermitian
552 ]
554 hams = [
555 self.hamiltonian_term(n, args=args, params=params)
556 for n in cell_terms
557 ]
558 # Add Hermitian conjugates
559 hams += [
560 ham.conjugate().transpose(0, 2, 1)
561 for ham, n in zip(hams, cell_terms)
562 if self.terms[n].hermitian
563 ]
565 _, norbs, orb_offsets = self.site_ranges.transpose()
567 shape = (orb_offsets[next_cell], orb_offsets[next_cell])
568 func = _vectorized_make_sparse if sparse else _vectorized_make_dense
569 mat = func(subgraphs, hams, norbs, orb_offsets, site_offsets, shape=shape)
571 return mat
574@deprecate_args
575@cython.binding(True)
576def vectorized_inter_cell_hopping(self, args=(), sparse=False, *, params=None):
577 """Hopping Hamiltonian between two cells of the infinite system.
579 This method returns a complex matrix that represents the hopping from
580 the *interface sites* of unit cell ``n - 1`` to *all* the sites of
581 unit cell ``n``. It is therefore generally a *rectangular* matrix of
582 shape ``(N_uc, N_iface)`` where ``N_uc`` is the number of orbitals
583 in the unit cell, and ``N_iface`` is the number of orbitals on the
584 *interface* sites (i.e. the sites with hoppings *to* the next unit cell).
586 Providing positional arguments via 'args' is deprecated,
587 instead, provide named parameters as a dictionary via 'params'.
588 """
590 site_offsets = np.cumsum([0] + [len(arr) for arr in self.site_arrays])
592 # This method is only meaningful for systems with a 1D translational
593 # symmetry, and we use this fact in several places
594 assert all(len(t.symmetry_element) == 1 for t in self.terms)
596 # Symmetry element -1 means hoppings *from* the *previous*
597 # unit cell. These are directly the hoppings we wish to return.
598 inter_cell_hopping_terms = [
599 n for n, t in enumerate(self.terms)
600 if t.symmetry_element[0] == -1
601 ]
602 # Symmetry element +1 means hoppings *from* the *next* unit cell.
603 # These are related by translational symmetry to hoppings *to*
604 # the *previous* unit cell. We therefore need the *reverse*
605 # (and conjugate) of these hoppings.
606 reversed_inter_cell_hopping_terms = [
607 n for n, t in enumerate(self.terms)
608 if t.symmetry_element[0] == +1
609 ]
611 inter_cell_hams = [
612 self.hamiltonian_term(n, args=args, params=params)
613 for n in inter_cell_hopping_terms
614 ]
615 reversed_inter_cell_hams = [
616 self.hamiltonian_term(n, args=args, params=params)
617 .conjugate().transpose(0, 2, 1)
618 for n in reversed_inter_cell_hopping_terms
619 ]
621 hams = inter_cell_hams + reversed_inter_cell_hams
623 subgraphs = [
624 self.subgraphs[self.terms[n].subgraph]
625 for n in inter_cell_hopping_terms
626 ]
627 subgraphs += [
628 _reverse_subgraph(self.subgraphs[self.terms[n].subgraph])
629 for n in reversed_inter_cell_hopping_terms
630 ]
632 _, norbs, orb_offsets = self.site_ranges.transpose()
634 # SiteArrays containing interface sites appear before SiteArrays
635 # containing non-interface sites, so the max of the site array
636 # indices that appear in the 'from' site arrays of the inter-cell
637 # hoppings allows us to get the number of interface orbitals.
638 last_iface_site_array = max(
639 from_site_array for (_, from_site_array), _ in subgraphs
640 )
641 iface_norbs = orb_offsets[last_iface_site_array + 1]
642 fd_norbs = orb_offsets[-1]
644 # TODO: return a square matrix when we no longer need to maintain
645 # backwards compatibility with unvectorized systems.
646 shape = (fd_norbs, iface_norbs)
647 func = _vectorized_make_sparse if sparse else _vectorized_make_dense
648 mat = func(subgraphs, hams, norbs, orb_offsets, site_offsets, shape=shape)
649 return mat