Coverage for kwant/operator.pyx : 84%

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
# # 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 # http://kwant-project.org/license. A list of Kwant authors can be found in # the file AUTHORS.rst at the top-level directory of this distribution and at # http://kwant-project.org/authors. """Tools for working with operators for acting on wavefunctions."""
import cython
from libc cimport math
from .graph.core cimport EdgeIterator from .graph.defs cimport gint
################ Generic Utility functions
@cython.boundscheck(False) @cython.wraparound(False) cdef gint _bisect(gint[:] a, gint x): "bisect.bisect specialized for searching `site_ranges`" else:
@cython.boundscheck(False) @cython.wraparound(False) cdef int _is_herm_conj(complex[:, :] a, complex[:, :] b, double atol=1e-300, double rtol=1e-13) except -1: "Return True if `a` is the Hermitian conjugate of `b`."
# compute max(a) cdef gint i, j
cdef complex ctmp
################ Helper functions
'the declared number of orbitals')
'`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." raise UserCodeError('Onsite matrix is not square') raise UserCodeError(_shape_msg.format('Onsite'))
cdef int _check_ham(complex[:, :] H, ham, args, params, gint a, gint a_norbs, gint b, gint b_norbs, int check_hermiticity) except -1: "Check Hamiltonian matrix for correct shape and hermiticity." raise UserCodeError(_shape_msg.format('Hamiltonian')) # call the "partner" element if we are not on the diagonal raise ValueError(_herm_msg.format('Hamiltonian'))
@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. # calculate the slice
@cython.boundscheck(False) @cython.wraparound(False)
cdef gint w, a, a_offset, a_norbs, b, b_offset, b_norbs else:
cdef gint _unused, tot_norbs
"""Normalize the format of `where` when `where` contains sites.
If `where` is None, then all sites in the system are returned. If it is a general iterator then it is expanded into an array. If `syst` is a finalized Builder then `where` should contain `Site` objects, otherwise it should contain integers. """ size = (syst.cell_size except AttributeError: _where = list(filter(where, range(syst.graph.num_nodes))) else: except AttributeError: _where = list(where) if any(w < 0 or w >= syst.graph.num_nodes for w in _where): raise ValueError('`where` contains sites that are not in the ' 'system.')
if any(w >= syst.cell_size for w in _where): raise ValueError('Only sites in the fundamental domain may be ' 'specified using `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 a finalized Builder then `where` should contain pairs of `Site` objects, otherwise it should contain pairs of integers. """ # we cannot extract the hoppings in the same order as they are in the # graph while simultaneously excluding all inter-cell hoppings raise ValueError('`where` must be provided when calculating ' 'current in an InfiniteSystem.') else: _where = list(filter(lambda h: where(*h), syst.graph)) else: except AttributeError: _where = list(where) # NOTE: if we ever have operators that contain elements that are # not in the system graph, then we should modify this check if any(not syst.graph.has_edge(*w) for w in where): raise ValueError('`where` contains hoppings that are not in the ' 'system.')
'using `where`.')
## These two classes are here to avoid using closures, as these will ## break pickle support. These are only used inside '_normalize_onsite'.
"""Normalize the format of `onsite`.
If `onsite` is a function or a mapping (dictionary) then a function is returned. """
# make 'onsite' compatible with hamiltonian value functions except AttributeError: _onsite = onsite raise TypeError('Provide `onsite` as a value or a function for ' 'systems that are not finalized Builders.')
# onsites known; immediately check for correct shape and hermiticity
else: # single onsite; immediately check for correct shape and hermiticity # 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 # we have the same number of orbitals everywhere 'but there are {1} orbitals per site in the system') else: 'different numbers of orbitals on different sites')
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)``. f : callable evaluates matrix blocks. Has signature ``(a, n_rows, b, n_cols)`` where all the arguments are integers and ``a`` and ``b`` are the contents of ``where``. This function must return a matrix of shape ``(n_rows, n_cols)``.
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. """
cdef public gint[:, :] block_offsets, block_shapes cdef public gint[:] data_offsets cdef public complex[:] data
@cython.embedsignature @cython.boundscheck(False) @cython.wraparound(False) def __init__(self, gint[:, :] where, gint[:, :] block_offsets, gint[:, :] block_shapes, f): raise ValueError('Arrays should be the same length along ' 'the first axis.') ### calculate shapes and data_offsets ### Populate data array cdef complex[:, :] mat cdef gint i, j, off, a, b, a_norbs, b_norbs # call the function that gives the matrix # Copy data
cdef complex* get(self, gint block_idx):
def __getstate__(self): )))
def __setstate__(self, state):
################ Local Observables
# supported operations within the `_operate` method ctypedef enum operation: MAT_ELS ACT
cdef class _LocalOperator: """Base class for operators defined by an on-site matrix and the Hamiltonian.
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``. 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]`). check_hermiticity : bool If True, checks that ``onsite``, as well as any relevant parts of the Hamiltonian are hermitian. sum : bool, default: False If True, then calling this operator will return a single scalar, otherwise a vector will be returned (see `~kwant.operator._LocalOperator.__call__` for details). """
cdef public int check_hermiticity, sum cdef public object syst, onsite, _onsite_params_info cdef public gint[:, :] where, _site_ranges cdef public BlockSparseMatrix _bound_onsite, _bound_hamiltonian
@cython.embedsignature def __init__(self, syst, onsite, where, *, check_hermiticity=True, sum=False): 'Declare the number of orbitals using the ' '`norbs` keyword argument when constructing ' 'the site families (lattices).')
@cython.embedsignature def __call__(self, bra, ket=None, args=(), *, params=None): r"""Return the matrix elements of the operator.
An operator ``A`` can be called like
>>> A(psi)
to compute the expectation value :math:`\bra{ψ} A \ket{ψ}`, or like
>>> A(phi, psi)
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.
Parameters ---------- bra, ket : sequence of complex 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. Mutually exclusive with 'params'. params : dict, optional Dictionary of parameter names and their values. Mutually exclusive with 'args'.
Returns ------- `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. """ raise ValueError("Extra arguments are already bound to this " "operator. You should call this operator " "providing neither 'args' nor 'params'.") raise TypeError('bra must be an array') raise ValueError('ket vector is incorrect shape')
# if `where` just contains sites, then we want a strictly 1D array
# if everything is Hermitian then result is real if bra == ket
@cython.embedsignature def act(self, ket, args=(), *, params=None): """Act with the operator on a wavefunction.
For an operator :math:`Q_{iαβ}` and ``ket`` :math:`ψ_β` this computes :math:`∑_{iβ} Q_{iαβ} ψ_β`.
Parameters ---------- ket : sequence of complex Wavefunctions defined over all the orbitals of the system. args : tuple The extra arguments to the Hamiltonian value functions and the operator ``onsite`` function. Mutually exclusive with 'params'. params : dict, optional Dictionary of parameter names and their values. Mutually exclusive with 'args'.
Returns ------- Array of `complex`. """ raise ValueError("Extra arguments are already bound to this " "operator. You should call this operator " "providing neither 'args' nor 'params'.")
raise TypeError('ket must be an array')
@cython.embedsignature def bind(self, args=(), *, params=None): """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. """ # generic creation of new instance # NOTE: subclasses should populate `bound_hamiltonian` if needed
def _operate(self, complex[:] out_data, complex[:] bra, complex[:] ket, args, operation op, *, params=None): """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 the operator ``onsite`` function. Mutually exclusive with 'params'. op : operation The operation to perform. `MAT_ELS`: calculate matrix elements between `bra` and `ket` `ACT`: act on `ket` with the operator params : dict, optional Dictionary of parameter names and their values. Mutually exclusive with 'args'. """ raise NotImplementedError()
cdef BlockSparseMatrix _eval_onsites(self, args, params): """Evaluate the onsite matrices on all elements of `where`"""
raise ValueError("Parameters {} have default values " "and may not be set with 'params'" .format(', '.join(invalid_params)))
cdef BlockSparseMatrix _eval_hamiltonian(self, args, params): """Evaluate the Hamiltonian on all elements of `where`."""
def __getstate__(self): )
def __setstate__(self, state):
cdef class Density(_LocalOperator): """An operator for calculating general densities.
An instance of this class can be called like a function to evaluate the expectation value with a wavefunction. See `~kwant.operator.Density.__call__` for details.
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. where : sequence of `int` or `~kwant.builder.Site`, or callable, optional 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, it should take a single `int` or `~kwant.builder.Site` (if ``syst`` is 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. sum : bool, default: False If True, then calling this operator will return a single scalar, otherwise a vector will be returned (see `~kwant.operator.Density.__call__` for details).
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. """
@cython.embedsignature def __init__(self, syst, onsite=1, where=None, *, check_hermiticity=True, sum=False):
@cython.boundscheck(False) @cython.wraparound(False) def _operate(self, complex[:] out_data, complex[:] bra, complex[:] ket, args, operation op, *, params=None): # prepare onsite matrices cdef complex[:, :] _tmp_mat cdef BlockSparseMatrix M_a_blocks
else:
# loop-local variables cdef gint a, a_s, a_norbs cdef gint i, j, w cdef complex tmp, bra_conj ### loop over sites ### get the next site, start orbital and number of orbitals ### get the next onsite matrix, if necessary ### do the actual calculation elif op == ACT:
@cython.boundscheck(False) @cython.wraparound(False) @cython.cdivision(True) @cython.embedsignature def tocoo(self, args=(), *, params=None): """Convert the operator to coordinate format sparse matrix.""" cdef int [:, :] offsets, shapes cdef int [:] row, col "operator. You should call this operator " "providing neither 'args' nor 'params'.")
else: else:
cdef class Current(_LocalOperator): r"""An operator for calculating general currents.
An instance of this class can be called like a function to evaluate the expectation value with a wavefunction. See `~kwant.operator.Current.__call__` for details.
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. where : sequence of pairs of `int` or `~kwant.builder.Site`, or callable, optional 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 `~kwant.builder.Site` (if ``syst`` is a finalized builder) and return 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. sum : bool, default: False If True, then calling this operator will return a single scalar, otherwise a vector will be returned (see `~kwant.operator.Current.__call__` for details).
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 :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.
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``. """
@cython.embedsignature def __init__(self, syst, onsite=1, where=None, *, check_hermiticity=True, sum=False):
@cython.embedsignature def bind(self, args=(), *, params=None): """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. """
@cython.boundscheck(False) @cython.wraparound(False) def _operate(self, complex[:] out_data, complex[:] bra, complex[:] ket, args, operation op, *, params=None): # prepare onsite matrices and hamiltonians cdef complex[:, :] _tmp_mat cdef BlockSparseMatrix M_a_blocks, H_ab_blocks
else:
else:
# main loop cdef gint a, a_s, a_norbs, b, b_s, b_norbs cdef gint i, j, k, w cdef complex tmp ### get the next hopping's start orbitals and numbers of orbitals ### get the next onsite and Hamiltonian matrices ### do the actual calculation H_ab[j * b_norbs + i].conjugate() * M_a[j * a_norbs + k] * elif op == ACT: 1j * H_ab[j * b_norbs + i].conjugate() * 1j * M_a[j * a_norbs + k] * H_ab[k * b_norbs + i] *
cdef class Source(_LocalOperator): """An operator for calculating general sources.
An instance of this class can be called like a function to evaluate the expectation value with a wavefunction. See `~kwant.operator.Source.__call__` for details.
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. where : sequence of `int` or `~kwant.builder.Site`, or callable, optional 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, it should take a single `int` or `~kwant.builder.Site` (if ``syst`` is 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. sum : bool, default: False If True, then calling this operator will return a single scalar, otherwise a vector will be returned (see `~kwant.operator.Source.__call__` for details).
Notes ----- An example of a "source" is a spin torque. 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_{i}` is the onsite Hamiltonian on site site `i`, then an instance of this class represents the tensor :math:`Q_{iαβ}` which is equal to :math:`(H_{i})^† M_i - M_i H_{i}` when α and β are orbitals on site :math:`i`, and zero otherwise. """
@cython.embedsignature def __init__(self, syst, onsite=1, where=None, *, check_hermiticity=True, sum=False):
@cython.embedsignature def bind(self, args=(), *, params=None): """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. """
@cython.boundscheck(False) @cython.wraparound(False) def _operate(self, complex[:] out_data, complex[:] bra, complex[:] ket, args, operation op, *, params=None): # prepare onsite matrices and hamiltonians cdef complex[:, :] _tmp_mat cdef BlockSparseMatrix M_a_blocks, H_aa_blocks
else:
else:
# main loop cdef gint a, a_s, a_norbs cdef gint i, j, k, w cdef complex tmp, tmp2 ### get the next site, start orbital and number of orbitals # row offsets and block size are the same as for columns, as # we are only dealing with the block-diagonal part of H ### get the next onsite and Hamiltonian matrices ### do the actual calculation - M_a[i * a_norbs + j] * elif op == ACT: - M_a[i * a_norbs + j] * |