operator.pyx 58 KB
Newer Older
1
# Copyright 2011-2019 Kwant authors.
Joseph Weston's avatar
Joseph Weston committed
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
5
# https://kwant-project.org/license.  A list of Kwant authors can be found in
Joseph Weston's avatar
Joseph Weston committed
6
# the file AUTHORS.rst at the top-level directory of this distribution and at
7
# https://kwant-project.org/authors.
8
"""Tools for working with operators for acting on wavefunctions."""
Joseph Weston's avatar
Joseph Weston committed
9
10
11
12

__all__ = ['Density', 'Current', 'Source']

import cython
13
import itertools
Joseph Weston's avatar
Joseph Weston committed
14
15
import functools as ft
import collections
16
import warnings
Anton Akhmerov's avatar
Anton Akhmerov committed
17

Joseph Weston's avatar
Joseph Weston committed
18
19
import numpy as np
import tinyarray as ta
Anton Akhmerov's avatar
Anton Akhmerov committed
20
from scipy.sparse import coo_matrix
Joseph Weston's avatar
Joseph Weston committed
21
22
23

from libc cimport math

24
25
26
cdef extern from "complex.h":
    double cabs(double complex)

Joseph Weston's avatar
Joseph Weston committed
27
from .graph.core cimport EdgeIterator
28
from .graph.core import DisabledFeatureError, NodeDoesNotExistError
Joseph Weston's avatar
Joseph Weston committed
29
30
from .graph.defs cimport gint
from .graph.defs import gint_dtype
31
32
33
from .system import (
    is_infinite, is_vectorized, Site, SiteArray, _normalize_matrix_blocks
)
34
35
36
from ._common import (
    UserCodeError, KwantDeprecationWarning, get_parameters, deprecate_args
)
Joseph Weston's avatar
Joseph Weston committed
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)
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's avatar
Joseph Weston committed
64
65
66

    # compute max(a)
    cdef double tmp, max_a = 0
67
    cdef gint i, j, k
Joseph Weston's avatar
Joseph Weston committed
68
69
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
70
            tmp = cabs(a[i, j])
Joseph Weston's avatar
Joseph Weston committed
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]):
77
78
        for j in range(i, a.shape[1]):
            tmp = cabs(a[i, j] - a[j, i].conjugate())
Joseph Weston's avatar
Joseph Weston committed
79
80
81
82
            if tmp > tol:
                return False
    return True

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's avatar
Joseph Weston committed
112

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's avatar
Joseph Weston committed
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'))
142
    if check_hermiticity and not _is_hermitian(M):
Joseph Weston's avatar
Joseph Weston committed
143
144
145
146
        raise ValueError(_herm_msg.format('Onsite'))
    return 0


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's avatar
Joseph Weston committed
162
        raise UserCodeError(_shape_msg.format('Hamiltonian'))
163
    if check_hermiticity and not _is_hermitian_3d(H):
Joseph Weston's avatar
Joseph Weston committed
164
165
166
167
            raise ValueError(_herm_msg.format('Hamiltonian'))
    return 0


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's avatar
Joseph Weston committed
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
219
    _check_norbs(syst)
220
    is_infinite_system = is_infinite(syst)
Joseph Weston's avatar
Joseph Weston committed
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.
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's avatar
Joseph Weston committed
233
234
235
    otherwise it should contain integers.
    """
    if where is None:
236
        if is_infinite(syst):
237
238
239
            where = list(range(syst.cell_size))
        else:
            where = list(range(syst.graph.num_nodes))
Joseph Weston's avatar
Joseph Weston committed
240
    elif callable(where):
241
        try:
242
            where = [syst.id_by_site[s] for s in filter(where, syst.sites)]
243
        except AttributeError:
244
            if is_infinite(syst):
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's avatar
Joseph Weston committed
248
    else:
249
        if isinstance(where[0], Site):
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's avatar
Joseph Weston committed
255

256
    where = np.asarray(where, dtype=gint_dtype).reshape(-1, 1)
Joseph Weston's avatar
Joseph Weston committed
257

258
    if is_infinite(syst) and np.any(where >= syst.cell_size):
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's avatar
Joseph Weston committed
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
273
    a finalized Builder then `where` may contain pairs of `Site` objects,
Joseph Weston's avatar
Joseph Weston committed
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
279
        if is_infinite(syst):
Joseph Weston's avatar
Joseph Weston committed
280
281
            raise ValueError('`where` must be provided when calculating '
                             'current in an InfiniteSystem.')
282
        where = list(syst.graph)
Joseph Weston's avatar
Joseph Weston committed
283
    elif callable(where):
284
        if hasattr(syst, "sites"):
285
            def idxwhere(hop):
Joseph Weston's avatar
Joseph Weston committed
286
287
                a, b = hop
                return where(syst.sites[a], syst.sites[b])
288
            where = list(filter(idxwhere, syst.graph))
Joseph Weston's avatar
Joseph Weston committed
289
        else:
290
            where = list(filter(lambda h: where(*h), syst.graph))
Joseph Weston's avatar
Joseph Weston committed
291
    else:
292
        if isinstance(where[0][0], Site):
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
301
        try:
302
303
            error = ValueError('`where` contains hoppings that are not '
                               'in the system.')
Joseph Weston's avatar
Joseph Weston committed
304
            if any(not syst.graph.has_edge(*w) for w in where):
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's avatar
Joseph Weston committed
311

312
    if is_infinite(syst) and np.any(where > syst.cell_size):
313
314
        raise ValueError('Only intra-cell hoppings may be specified '
                         'using `where`.')
Joseph Weston's avatar
Joseph Weston committed
315

316
    return where
Joseph Weston's avatar
Joseph Weston committed
317
318


319
## These four classes are here to avoid using closures, as these will
320
321
## break pickle support. These are only used inside '_normalize_onsite'.

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


332
333
class _FunctionalOnsite:

334
    def __init__(self, onsite, sites, site_ranges):
335
336
        self.onsite = onsite
        self.sites = sites
337
338
339
340
        self.site_ranges = site_ranges

    def __call__(self, site_range, site_offsets, *args):
        sites = self.sites
341
        offset, norbs, _ = self.site_ranges[site_range]
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)
346
347

        expected_shape = (len(site_offsets), norbs, norbs)
348
349
        return _normalize_matrix_blocks(ret, expected_shape,
                                        calling_function=self.onsite)
350

351

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)
366
367

        expected_shape = (len(sites), sites.family.norbs, sites.family.norbs)
368
369
        return _normalize_matrix_blocks(ret, expected_shape,
                                        calling_function=self.onsite)
370
371


372
class _FunctionalOnsiteNoTransform:
373

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):
379
380
        offset, norbs, _ = self.site_ranges[site_range]
        site_ids = offset + site_offsets
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)
385
386

        expected_shape = (len(site_offsets), norbs, norbs)
387
388
        return _normalize_matrix_blocks(ret, expected_shape,
                                        calling_function=self.onsite)
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)
400
401
402

        expected_shape = (len(site_offsets), fam.norbs, fam.norbs)
        return _normalize_matrix_blocks(ret, expected_shape)
403
404


Joseph Weston's avatar
Joseph Weston committed
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.
    """
411
    param_names = ()
412

Joseph Weston's avatar
Joseph Weston committed
413
    if callable(onsite):
414
        # make 'onsite' compatible with hamiltonian value functions
415
        param_names = get_parameters(onsite)[1:]
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's avatar
Joseph Weston committed
422

423
    elif isinstance(onsite, collections.abc.Mapping):
Joseph Weston's avatar
Joseph Weston committed
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)

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's avatar
Joseph Weston committed
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)
449
        elif len(set(norbs for _, norbs, _ in syst.site_ranges[:-1])) == 1:
Joseph Weston's avatar
Joseph Weston committed
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)

461
    return _onsite, param_names
Joseph Weston's avatar
Joseph Weston committed
462
463


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]
497
498
        ((term_sa, _), (term_site_offsets, _)) = syst.subgraphs[term.subgraph]
        term_sites = site_offsets[term_sa] + term_site_offsets
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's avatar
Joseph Weston committed
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)``.
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's avatar
Joseph Weston committed
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,
595
                  gint[:, :] block_shapes, matrix_elements):
Joseph Weston's avatar
Joseph Weston committed
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)
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's avatar
Joseph Weston committed
623
624
625
626

    cdef complex* get(self, gint block_idx):
        return  <complex*> &self.data[0] + self.data_offsets[block_idx]

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's avatar
Joseph Weston committed
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:
652
653
    """Base class for operators defined by an on-site matrix and the
    Hamiltonian.
Joseph Weston's avatar
Joseph Weston committed
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``.
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's avatar
Joseph Weston committed
674
675
676
    check_hermiticity : bool
        If True, checks that ``onsite``, as well as any relevant parts
        of the Hamiltonian are hermitian.
677
    sum : bool, default: False
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's avatar
Joseph Weston committed
681
682
683
    """

    @cython.embedsignature
684
685
    def __init__(self, syst, onsite, where, *,
                 check_hermiticity=True, sum=False):
686
        _check_norbs(syst)
687
688
689
        if is_vectorized(syst) and is_infinite(syst):
            raise TypeError('Vectorized infinite systems cannot yet be '
                            'used with operators.')
Joseph Weston's avatar
Joseph Weston committed
690
691

        self.syst = syst
692
693
        self.onsite, self._onsite_param_names = _normalize_onsite(
            syst, onsite, check_hermiticity)
Joseph Weston's avatar
Joseph Weston committed
694
        self.check_hermiticity = check_hermiticity
695
        self.sum = sum
Joseph Weston's avatar
Joseph Weston committed
696
        self._site_ranges = np.asarray(syst.site_ranges, dtype=gint_dtype)
697
        self.where = where
Joseph Weston's avatar
Joseph Weston committed
698
699
700
        self._bound_onsite = None
        self._bound_hamiltonian = None

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's avatar
Joseph Weston committed
735
    @cython.embedsignature
736
737
    def __call__(self, bra, ket=None, args=(), *, params=None):
        r"""Return the matrix elements of the operator.
Joseph Weston's avatar
Joseph Weston committed
738

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)

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.
769

Joseph Weston's avatar
Joseph Weston committed
770
771
        Parameters
        ----------
772
        bra, ket : sequence of complex
Joseph Weston's avatar
Joseph Weston committed
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.
779
            Deprecated in favor of 'params' (and mutually exclusive with it).
780
781
782
        params : dict, optional
            Dictionary of parameter names and their values. Mutually exclusive
            with 'args'.
Joseph Weston's avatar
Joseph Weston committed
783
784
785

        Returns
        -------
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's avatar
Joseph Weston committed
789
        """
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'.")
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()
799
800
        if args and params:
            raise TypeError("'args' and 'params' are mutually exclusive.")
Joseph Weston's avatar
Joseph Weston committed
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'
808
            msg = 'bra ' + msg if ket is None else msg
Joseph Weston's avatar
Joseph Weston committed
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)
814
815
        self._operate(out_data=result, bra=bra, ket=ket, args=args,
                      params=params, op=MAT_ELS)
816
817
        # if everything is Hermitian then result is real if bra == ket
        if self.check_hermiticity and bra is ket:
Joseph Weston's avatar
Joseph Weston committed
818
            result = result.real
819
        return np.sum(result) if self.sum else result
Joseph Weston's avatar
Joseph Weston committed
820
821

    @cython.embedsignature
822
    def act(self, ket, args=(), *, params=None):
Joseph Weston's avatar
Joseph Weston committed
823
824
        """Act with the operator on a wavefunction.

825
826
827
        For an operator :math:`Q_{iαβ}` and ``ket`` :math:`ψ_β`
        this computes :math:`∑_{iβ} Q_{iαβ} ψ_β`.

Joseph Weston's avatar
Joseph Weston committed
828
829
        Parameters
        ----------
830
        ket : sequence of complex
Joseph Weston's avatar
Joseph Weston committed
831
832
833
            Wavefunctions defined over all the orbitals of the system.
        args : tuple
            The extra arguments to the Hamiltonian value functions and
834
835
            the operator ``onsite`` function.
            Deprecated in favor of 'params' (and mutually exclusive with it).
836
837
838
        params : dict, optional
            Dictionary of parameter names and their values. Mutually exclusive
            with 'args'.
Joseph Weston's avatar
Joseph Weston committed
839
840
841

        Returns
        -------
842
        Array of `complex`.
Joseph Weston's avatar
Joseph Weston committed
843
        """
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'.")
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()
853
854
        if args and params:
            raise TypeError("'args' and 'params' are mutually exclusive.")
Joseph Weston's avatar
Joseph Weston committed
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')
862
        result = np.zeros((tot_norbs,), dtype=complex)
863
864
        self._operate(out_data=result, bra=None, ket=ket, args=args,
                      params=params, op=ACT)
Joseph Weston's avatar
Joseph Weston committed
865
866
867
        return result

    @cython.embedsignature
868
    def bind(self, args=(), *, params=None):
Joseph Weston's avatar
Joseph Weston committed
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.
873
874
875

        Providing positional arguments via 'args' is deprecated,
        instead provide named parameters as a dictionary via 'params'.
Joseph Weston's avatar
Joseph Weston committed
876
        """
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()
882
883
        if args and params:
            raise TypeError("'args' and 'params' are mutually exclusive.")
Joseph Weston's avatar
Joseph Weston committed
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
889
        q._onsite_param_names = self._onsite_param_names
Joseph Weston's avatar
Joseph Weston committed
890
        q.where = self.where
891
        q.sum = self.sum
Joseph Weston's avatar
Joseph Weston committed
892
893
894
        q._site_ranges = self._site_ranges
        q.check_hermiticity = self.check_hermiticity
        if callable(self.onsite):
895
            q._bound_onsite = self._eval_onsites(args, params)
Joseph Weston's avatar
Joseph Weston committed
896
897
898
899
        # NOTE: subclasses should populate `bound_hamiltonian` if needed
        return q

    def _operate(self, complex[:] out_data, complex[:] bra, complex[:] ket,
900
                 args, operation op, *, params=None):
Joseph Weston's avatar
Joseph Weston committed
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
914
915
            the operator ``onsite`` function.
            Deprecated in favor of 'params' (and mutually exclusive with it).
Joseph Weston's avatar
Joseph Weston committed
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
920
921
922
        params : dict, optional
            Dictionary of parameter names and their values. Mutually exclusive
            with 'args'.
Joseph Weston's avatar
Joseph Weston committed
923
924
925
        """
        raise NotImplementedError()

926
    cdef BlockSparseMatrix _eval_onsites(self, args, params):
Joseph Weston's avatar
Joseph Weston committed
927
928
        """Evaluate the onsite matrices on all elements of `where`"""
        assert callable(self.onsite)
929
        assert not (args and params)
Joseph Weston's avatar
Joseph Weston committed
930
        check_hermiticity = self.check_hermiticity
931
932
933
        syst = self.syst

        _is_vectorized = is_vectorized(syst)
Joseph Weston's avatar
Joseph Weston committed
934

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))
944

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)
959
            _check_onsites(data, norbs, self.check_hermiticity)
960
961
962
            return data

        matrix_elements = _make_matrix_elements(eval_onsite, self._terms)
Joseph Weston's avatar
Joseph Weston committed
963
        offsets, norbs = _get_all_orbs(self.where, self._site_ranges)
964
965
        return  BlockSparseMatrix(self.where, offsets, norbs, matrix_elements)

Joseph Weston's avatar
Joseph Weston committed
966

967
    cdef BlockSparseMatrix _eval_hamiltonian(self, args, params):
Joseph Weston's avatar
Joseph Weston committed
968
        """Evaluate the Hamiltonian on all elements of `where`."""
969
970
971
972

        where = self.where
        syst = self.syst
        is_onsite = self.where.shape[1] == 1
Joseph Weston's avatar
Joseph Weston committed
973
974
        check_hermiticity = self.check_hermiticity

975
        if is_vectorized(self.syst):
Joseph Weston's avatar
Joseph Weston committed
976

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)
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)
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
                    ]
1013

1014
                (to_sr, from_sr) = term_id
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

1020
                _check_hams(data, to_norbs, from_norbs, is_onsite and check_hermiticity)
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's avatar
Joseph Weston committed
1027

1028
1029
1030
    def __getstate__(self):
        return (
            (self.check_hermiticity, self.sum),
1031
            (self.syst, self.onsite, self._onsite_param_names),
1032
            tuple(map(np.asarray, (self.where, self._site_ranges))),
1033
            (self._terms,),
1034
1035
1036
1037
1038
            (self._bound_onsite, self._bound_hamiltonian),
        )

    def __setstate__(self, state):
        ((self.check_hermiticity, self.sum),
1039
         (self.syst, self.onsite, self._onsite_param_names),
1040
         (self.where, self._site_ranges),
1041
         (self._terms,),
1042
1043
1044
         (self._bound_onsite, self._bound_hamiltonian),
        ) = state

Joseph Weston's avatar
Joseph Weston committed
1045
1046
1047
1048

cdef class Density(_LocalOperator):
    """An operator for calculating general densities.

1049
    An instance of this class can be called like a function to evaluate the
1050
1051
    expectation value with a wavefunction. See
    `~kwant.operator.Density.__call__` for details.
Joseph Weston's avatar
Joseph Weston committed
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.
1061
    where : sequence of `int` or `~kwant.system.Site`, or callable, optional
Joseph Weston's avatar
Joseph Weston committed
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,
1064
        it should take a single `int` or `~kwant.system.Site` (if ``syst`` is
Joseph Weston's avatar
Joseph Weston committed
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.
1071
    sum : bool, default: False
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's avatar
Joseph Weston committed
1075

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's avatar
Joseph Weston committed
1083
1084
1085
    """

    @cython.embedsignature
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,
1090
                         check_hermiticity=check_hermiticity, sum=sum)
Joseph Weston's avatar
Joseph Weston committed
1091
1092
1093
1094

    @cython.boundscheck(False)
    @cython.wraparound(False)
    def _operate(self, complex[:] out_data, complex[:] bra, complex[:] ket,
1095
                 args, operation op, *, params=None):
Joseph Weston's avatar
Joseph Weston committed
1096
1097
1098
1099
        matrix = ta.matrix
        cdef int unique_onsite = not callable(self.onsite)
        # prepare onsite matrices
        cdef complex[:, :] _tmp_mat
1100
        cdef complex *M_a = NULL
Joseph Weston's avatar
Joseph Weston committed
1101
        cdef BlockSparseMatrix M_a_blocks
1102

Joseph Weston's avatar
Joseph Weston committed
1103
1104
1105
        if unique_onsite:
            _tmp_mat = self.onsite
            M_a = <complex*> &_tmp_mat[0, 0]
1106
1107
        elif self._bound_onsite:
            M_a_blocks = self._bound_onsite
Joseph Weston's avatar
Joseph Weston committed
1108
        else:
1109
            M_a_blocks = self._eval_onsites(args, params)
1110

Joseph Weston's avatar
Joseph Weston committed
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's avatar
Anton Akhmerov committed
1138
1139
1140
    @cython.boundscheck(False)
    @cython.wraparound(False)
    @cython.cdivision(True)
1141
    @cython.embedsignature
1142
    def tocoo(self, args=(), *, params=None):
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's avatar
Anton Akhmerov committed
1148
1149
1150
        cdef int blk, blk_size, n_blocks, n, k = 0
        cdef int [:, :] offsets, shapes
        cdef int [:] row, col
1151
        if self._bound_onsite and (args or params):
Anton Akhmerov's avatar
Anton Akhmerov committed
1152
1153
           raise ValueError("Extra arguments are already bound to this "
                            "operator. You should call this operator "
1154
                            "providing neither 'args' nor 'params'.")
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()
1160
1161
        if args and params:
            raise TypeError("'args' and 'params' are mutually exclusive.")
Anton Akhmerov's avatar
Anton Akhmerov committed
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:
1171
1172
1173
            if self._bound_onsite is not None:
                onsite_matrix = self._bound_onsite
            else:
1174
                onsite_matrix = self._eval_onsites(args, params)
Anton Akhmerov's avatar
Anton Akhmerov committed
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's avatar
Joseph Weston committed
1193
1194

cdef class Current(_LocalOperator):
1195
    r"""An operator for calculating general currents.
Joseph Weston's avatar
Joseph Weston committed
1196

1197
    An instance of this class can be called like a function to evaluate the
1198
1199
    expectation value with a wavefunction. See
    `~kwant.operator.Current.__call__` for details.
Joseph Weston's avatar
Joseph Weston committed
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.
1210
    where : sequence of pairs of `int` or `~kwant.system.Site`, or callable, optional
Joseph Weston's avatar
Joseph Weston committed
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
1214
        `~kwant.system.Site` (if ``syst`` is a finalized builder) and return
Joseph Weston's avatar
Joseph Weston committed
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.
1221
    sum : bool, default: False
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's avatar
Joseph Weston committed
1225

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
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.
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's avatar
Joseph Weston committed
1238
1239
1240
    """

    @cython.embedsignature
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,
1245
                         check_hermiticity=check_hermiticity, sum=sum)
Joseph Weston's avatar
Joseph Weston committed
1246
1247

    @cython.embedsignature
1248
    def bind(self, args=(), *, params=None):
Joseph Weston's avatar
Joseph Weston committed
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.
1253
1254
1255

        Providing positional arguments via 'args' is deprecated,
        instead provide named parameters as a dictionary via 'params'.
Joseph Weston's avatar
Joseph Weston committed
1256
        """
1257
1258
        q = super().bind(args, params=params)
        q._bound_hamiltonian = self._eval_hamiltonian(args, params)
Joseph Weston's avatar
Joseph Weston committed
1259
1260
1261
1262
1263
        return q

    @cython.boundscheck(False)
    @cython.wraparound(False)
    def _operate(self, complex[:] out_data, complex[:] bra, complex[:] ket,
1264
                 args, operation op, *, params=None):
Joseph Weston's avatar
Joseph Weston committed
1265
1266
1267
        # prepare onsite matrices and hamiltonians
        cdef int unique_onsite = not callable(self.onsite)
        cdef complex[:, :] _tmp_mat
1268
1269
        cdef complex *M_a = NULL
        cdef complex *H_ab = NULL
Joseph Weston's avatar
Joseph Weston committed
1270
        cdef BlockSparseMatrix M_a_blocks, H_ab_blocks
1271

Joseph Weston's avatar
Joseph Weston committed
1272
1273
1274
        if unique_onsite:
            _tmp_mat = self.onsite
            M_a = <complex*> &_tmp_mat[0, 0]
1275
1276
1277
        elif self._bound_onsite:
            M_a_blocks = self._bound_onsite
        else:
1278
            M_a_blocks = self._eval_onsites(args, params)
1279
1280
1281

        if self._bound_hamiltonian:
            H_ab_blocks = self._bound_hamiltonian
Joseph Weston's avatar
Joseph Weston committed
1282
        else:
1283
            H_ab_blocks = self._eval_hamiltonian(args, params)
1284

Joseph Weston's avatar
Joseph Weston committed
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.

1329
    An instance of this class can be called like a function to evaluate the
1330
1331
    expectation value with a wavefunction. See
    `~kwant.operator.Source.__call__` for details.
Joseph Weston's avatar
Joseph Weston committed
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.
1342
    where : sequence of `int` or `~kwant.system.Site`, or callable, optional
Joseph Weston's avatar
Joseph Weston committed
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,
1345
        it should take a single `int` or `~kwant.system.Site` (if ``syst`` is
Joseph Weston's avatar
Joseph Weston committed
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.
1352
    sum : bool, default: False
1353