diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 7af2fa4c7f71374c9bc77861617e0a0b78b9d0ce..c6e2f7921e43b6bb5997f11e2889f91b46927071 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -62,6 +62,7 @@ build documentation:
run tests:
stage: test
script:
+ - pip3 install sympy
- py.test --cov=kwant --cov-report term --cov-report html --flakes kwant
artifacts:
paths:
diff --git a/kwant/__init__.py b/kwant/__init__.py
index 6dcc99eec2062e2befca40811efefa408e161c8e..a3c430d765b53497f8527c2b1485ca90e4bdbf65 100644
--- a/kwant/__init__.py
+++ b/kwant/__init__.py
@@ -9,6 +9,7 @@
__all__ = []
import numpy # Needed by C. Gohlke's Windows package.
+import warnings
try:
from . import _system
@@ -62,3 +63,12 @@ def test(verbose=True):
"-s"] + (['-v'] if verbose else []))
test.__test__ = False
+
+
+# Exposing discretizer
+try:
+ from .continuum import discretize
+ __all__.extend(['discretize'])
+except ImportError:
+ warnings.warn('Discretizer module not available. Is sympy installed?',
+ RuntimeWarning)
diff --git a/kwant/continuum/__init__.py b/kwant/continuum/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..bcc0b6d17d732799617f44c0f399592e4b10860a
--- /dev/null
+++ b/kwant/continuum/__init__.py
@@ -0,0 +1,7 @@
+__all__ = ['discretize', 'discretize_symbolic', 'build_discretized',
+ 'momentum_operators', 'position_operators']
+
+
+from .discretizer import discretize, discretize_symbolic, build_discretized
+from ._common import momentum_operators, position_operators
+from ._common import sympify, lambdify, make_commutative
diff --git a/kwant/continuum/_common.py b/kwant/continuum/_common.py
new file mode 100644
index 0000000000000000000000000000000000000000..b107d59f855d57373fc7bec027dee2c254eacbce
--- /dev/null
+++ b/kwant/continuum/_common.py
@@ -0,0 +1,217 @@
+import numpy as np
+
+import sympy
+from sympy.core.function import AppliedUndef
+from sympy.core.sympify import converter
+from sympy.abc import _clash
+
+import functools
+import inspect
+
+from collections import defaultdict
+from operator import mul
+
+from sympy.physics.matrices import msigma as _msigma
+from sympy.physics.quantum import TensorProduct
+
+
+momentum_operators = sympy.symbols('k_x k_y k_z', commutative=False)
+position_operators = sympy.symbols('x y z', commutative=False)
+
+msigma = lambda i: sympy.eye(2) if i == 0 else _msigma(i)
+pauli = (msigma(0), msigma(1), msigma(2), msigma(3))
+
+_clash = _clash.copy()
+_clash.update({s.name: s for s in momentum_operators})
+_clash.update({s.name: s for s in position_operators})
+_clash.update({'kron': TensorProduct, 'eye': sympy.eye})
+_clash.update({'sigma_{}'.format(c): p for c, p in zip('0xyz', pauli)})
+
+
+# workaroud for https://github.com/sympy/sympy/issues/12060
+del _clash['I']
+del _clash['pi']
+
+
+################ Various helpers to handle sympy
+
+
+def lambdify(hamiltonian, *, substitutions=None):
+ """Return a callable object for computing continuum Hamiltonian.
+
+ Parameters
+ ----------
+ hamiltonian : sympy.Expr or sympy.Matrix, or string
+ Symbolic representation of a continous Hamiltonian. When providing
+ a sympy expression, it is recommended to use operators
+ defined in `~kwant.continuum.momentum_operators` in order to provide
+ proper commutation properties. If a string is provided it will be
+ converted to a sympy expression using `kwant.continuum.sympify`.
+ substitutions : dict, defaults to empty
+ A namespace to be passed to ``kwant.continuum.sympify`` when
+ ``hamiltonian`` is a string. Could be used to simplify matrix input:
+ ``sympify('k_x**2 * s_z', substitutions={'s_z': [[1, 0], [0, -1]]})``.
+ """
+ expr = hamiltonian
+ if not isinstance(expr, (sympy.Expr, sympy.matrices.MatrixBase)):
+ expr = sympify(expr, substitutions)
+
+ args = [s.name for s in expr.atoms(sympy.Symbol)]
+ args += [str(f.func) for f in expr.atoms(AppliedUndef, sympy.Function)]
+
+ f = sympy.lambdify(sorted(args), expr)
+
+ sig = inspect.signature(f)
+ pars = list(sig.parameters.values())
+ pars = [p.replace(kind=inspect.Parameter.KEYWORD_ONLY) for p in pars]
+ f.__signature__ = inspect.Signature(pars)
+ return f
+
+
+def sympify(string, substitutions=None, **kwargs):
+ """Return sympified object with respect to kwant-specific rules.
+
+ This is a modification of ``sympy.sympify`` to apply kwant-specific rules,
+ which includes preservation of proper commutation relations between
+ position and momentum operators, and array to matrix casting.
+
+ Parameters
+ ----------
+ string : string
+ String representation of a Hamiltonian. Momenta must be defined as
+ ``k_i`` where ``i`` stands for ``x``, ``y``, or ``z``. All present
+ momenta and coordinates will be interpreted as non commutative.
+ substitutions : dict, defaults to empty
+ A namespace to be passed internally to ``sympy.sympify``.
+ Could be used to simplify matrix input:
+ ``sympify('k_x**2 * s_z', substitutions={'s_z': [[1, 0], [0, -1]]})``.
+ **kwargs
+ Additional arguments that will be passed to ``sympy.sympify``.
+
+ Example
+ -------
+ >>> from kwant.continuum import sympify
+ >>> hamiltonian = sympify('k_x * A(x) * k_x + V(x)')
+ """
+ stored_value = None
+
+ if substitutions is None:
+ substitutions = {}
+
+ substitutions.update(_clash)
+
+ try:
+ stored_value = converter.pop(list, None)
+ converter[list] = lambda x: sympy.Matrix(x)
+ substitutions = {k: (sympy.sympify(v, locals=substitutions, **kwargs)
+ if isinstance(v, (list, str)) else v)
+ for k, v in substitutions.items()}
+ hamiltonian = sympy.sympify(string, locals=substitutions, **kwargs)
+ hamiltonian = sympy.sympify(hamiltonian, **kwargs)
+ finally:
+ if stored_value is not None:
+ converter[list] = stored_value
+ else:
+ del converter[list]
+ return sympy.expand(hamiltonian)
+
+
+def make_commutative(expr, *symbols):
+ """Make sure that specified symbols are defined as commutative.
+
+ Parameters
+ ----------
+ expr: sympy.Expr or sympy.Matrix
+ symbols: sequace of symbols
+ Set of symbols that are requiered to be commutative. It doesn't matter
+ of symbol is provided as commutative or not.
+
+ Returns
+ -------
+ input expression with all specified symbols changed to commutative.
+ """
+ symbols = [sympy.Symbol(s.name, commutative=False) for s in symbols]
+ subs = {s: sympy.Symbol(s.name) for s in symbols}
+ expr = expr.subs(subs)
+ return expr
+
+
+def expression_monomials(expression, *gens):
+ """Parse ``expression`` into monomials in the symbols in ``gens``.
+
+ Example
+ -------
+ >>> expr = A * (x**2 + y) + B * x + C
+ >>> expression_monomials(expr, x, y)
+ {1: C, x**2: A, y: A, x: B}
+ """
+ if expression.atoms(AppliedUndef):
+ raise NotImplementedError('Getting monomials of expressions containing '
+ 'functions is not implemented.')
+
+ expression = make_commutative(expression, *gens)
+ gens = [make_commutative(g, g) for g in gens]
+
+ expression = sympy.expand(expression)
+ summands = expression.as_ordered_terms()
+
+ output = defaultdict(int)
+ for summand in summands:
+ key = [sympy.Integer(1)]
+ if summand in gens:
+ key.append(summand)
+
+ elif isinstance(summand, sympy.Pow):
+ if summand.args[0] in gens:
+ key.append(summand)
+
+ else:
+ for arg in summand.args:
+ if arg in gens:
+ key.append(arg)
+ if isinstance(arg, sympy.Pow):
+ if arg.args[0] in gens:
+ key.append(arg)
+
+ key = functools.reduce(mul, key)
+ val = summand.xreplace({g: 1 for g in gens})
+
+ ### to not create key
+ if val != 0:
+ output[key] += val
+
+ new_expression = sum(k * v for k, v in output.items())
+ assert sympy.expand(expression) == sympy.expand(new_expression)
+
+ return dict(output)
+
+
+def matrix_monomials(matrix, *gens):
+ output = defaultdict(lambda: sympy.zeros(*matrix.shape))
+ for (i, j), expression in np.ndenumerate(matrix):
+ summands = expression_monomials(expression, *gens)
+ for key, val in summands.items():
+ output[key][i, j] += val
+
+ return dict(output)
+
+
+################ general help functions
+
+def gcd(*args):
+ if len(args) == 1:
+ return args[0]
+
+ L = list(args)
+
+ while len(L) > 1:
+ a = L[len(L) - 2]
+ b = L[len(L) - 1]
+ L = L[:len(L) - 2]
+
+ while a:
+ a, b = b%a, a
+
+ L.append(b)
+
+ return abs(b)
diff --git a/kwant/continuum/discretizer.py b/kwant/continuum/discretizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..af7cd1c8831920e22a5ee10bb62d708484705086
--- /dev/null
+++ b/kwant/continuum/discretizer.py
@@ -0,0 +1,555 @@
+from collections import defaultdict
+
+import numpy as np
+import tinyarray as ta
+import sympy
+
+from sympy.utilities.lambdify import lambdastr
+from sympy.printing.lambdarepr import LambdaPrinter
+from sympy.printing.precedence import precedence
+
+from sympy.core.function import AppliedUndef
+
+from ..builder import Builder, HoppingKind
+from ..lattice import Monatomic, TranslationalSymmetry
+
+from ._common import sympify, gcd
+from ._common import position_operators, momentum_operators
+from ._common import matrix_monomials
+
+
+################ Globals variables and definitions
+
+_wf = sympy.Function('_internal_unique_name', commutative=False)
+_position_operators = {s.name: s for s in position_operators}
+_displacements = {s: sympy.Symbol('_internal_a_{}'.format(s)) for s in 'xyz'}
+
+
+################ Interface functions
+
+def discretize(hamiltonian, discrete_coordinates=None, lattice_constant=1,
+ substitutions=None, verbose=False):
+ """Construct a tight-binding model from a continuum Hamiltonian.
+
+ Parameters
+ ----------
+ hamiltonian : sympy.Expr or sympy.Matrix, or string
+ Symbolic representation of a continous Hamiltonian. When providing
+ a sympy expression, it is recommended to use operators
+ defined in `~kwant.continuum.momentum_operators` in order to provide
+ proper commutation properties. If a string is provided it will be
+ converted to a sympy expression using `kwant.continuum.sympify`.
+ discrete_coordinates : sequence of strings, or ``None`` (default)
+ Set of coordinates for which momentum operators will be treated as
+ differential operators. For example ``discrete_coordinates=('x', 'y')``.
+ If not provided they will be obtained from the input hamiltonian by
+ reading present coordinates and momentum operators. Order of discrete
+ coordinates is always lexical, even if provided otherwise.
+ lattice_constant : int or float, default: 1
+ Lattice constant for the template Builder.
+ substitutions : dict, defaults to empty
+ A namespace to be passed to `kwant.continuum.sympify` when
+ ``hamiltonian`` is a string. Could be used to simplify matrix input:
+ ``sympify('k_x**2 * s_z', substitutions={'s_z': [[1, 0], [0, -1]]})``.
+ verbose : bool, default: False
+ If ``True`` additional information will be printed.
+
+ Returns
+ -------
+ `~kwant.builder.Builder` with translational symmetry,
+ which can be used as a template.
+ """
+
+ args = hamiltonian, discrete_coordinates, substitutions, verbose
+ tb, coords = discretize_symbolic(*args)
+ return build_discretized(tb, coords, lattice_constant, substitutions, verbose)
+
+
+def discretize_symbolic(hamiltonian, discrete_coordinates=None, substitutions=None,
+ verbose=False):
+ """Discretize a continuous Hamiltonian into a tight-binding representation.
+
+ Parameters
+ ----------
+ hamiltonian : sympy.Expr or sympy.Matrix, or string
+ Symbolic representation of a continous Hamiltonian. When providing
+ a sympy expression, it is recommended to use operators
+ defined in `~kwant.continuum.momentum_operators` in order to provide
+ proper commutation properties. If a string is provided it will be
+ converted to a sympy expression using `kwant.continuum.sympify`.
+ discrete_coordinates : sequence of strings, or ``None`` (default)
+ Set of coordinates for which momentum operators will be treated as
+ differential operators. For example ``discrete_coordinates=('x', 'y')``.
+ If not provided they will be obtained from the input hamiltonian by
+ reading present coordinates and momentum operators. Order of discrete
+ coordinates is always lexical, even if provided otherwise.
+ substitutions : dict, defaults to empty
+ A namespace to be passed to `kwant.continuum.sympify` when
+ ``hamiltonian`` is a string. Could be used to simplify matrix input:
+ ``sympify('k_x**2 * s_z', substitutions={'s_z': [[1, 0], [0, -1]]})``.
+ verbose : bool, default: False
+ If ``True`` additional information will be printed.
+
+ Returns
+ -------
+ (discrete_hamiltonian, discrete_coordinates)
+ discrete_hamiltonian: dict
+ Keys are tuples of integers; the offsets of the hoppings
+ ((0, 0, 0) for the onsite). Values are symbolic expressions
+ for the hoppings/onsite.
+ discrete_coordinates : sequence of strings
+ The coordinates that have been discretized.
+ """
+ if not isinstance(hamiltonian, (sympy.Expr, sympy.matrices.MatrixBase)):
+ hamiltonian = sympify(hamiltonian, substitutions)
+
+ atoms = hamiltonian.atoms(sympy.Symbol)
+ if not all('a' != s.name for s in atoms):
+ raise ValueError("'a' is a symbol used internally to represent "
+ "lattice spacing; please use a different symbol.")
+
+ hamiltonian = sympy.expand(hamiltonian)
+ if discrete_coordinates is None:
+ used_momenta = set(momentum_operators) & set(atoms)
+ discrete_coordinates = {k.name[-1] for k in used_momenta}
+ else:
+ discrete_coordinates = set(discrete_coordinates)
+ if not discrete_coordinates <= set('xyz'):
+ msg = "Discrete coordinates must only contain 'x', 'y', or 'z'."
+ raise ValueError(msg)
+
+ discrete_coordinates = sorted(discrete_coordinates)
+
+ if len(discrete_coordinates) == 0:
+ raise ValueError("Failed to read any discrete coordinates. This is "
+ "probably due to a lack of momentum operators in "
+ "your input. You can use the 'discrete_coordinates'"
+ "parameter to provide them.")
+
+ if verbose:
+ print('Discrete coordinates set to: ',
+ discrete_coordinates, end='\n\n')
+
+ onsite_zeros = (0,) * len(discrete_coordinates)
+
+ if not isinstance(hamiltonian, sympy.matrices.MatrixBase):
+ hamiltonian = sympy.Matrix([hamiltonian])
+ _input_format = 'expression'
+ else:
+ _input_format = 'matrix'
+
+ shape = hamiltonian.shape
+ tb = defaultdict(lambda: sympy.zeros(*shape))
+ tb[onsite_zeros] = sympy.zeros(*shape)
+
+ for (i, j), expression in np.ndenumerate(hamiltonian):
+ hoppings = _discretize_expression(expression, discrete_coordinates)
+
+ for offset, hop in hoppings.items():
+ tb[offset][i, j] += hop
+
+ # do not include Hermitian conjugates of hoppings
+ wanted_hoppings = sorted(list(tb))[len(list(tb)) // 2:]
+ tb = {k: v for k, v in tb.items() if k in wanted_hoppings}
+
+ if _input_format == 'expression':
+ tb = {k: v[0, 0] for k, v in tb.items()}
+
+ return tb, discrete_coordinates
+
+
+def build_discretized(tb_hamiltonian, discrete_coordinates,
+ lattice_constant=1, substitutions=None, verbose=False):
+ """Create a template Builder from a symbolic tight-binding Hamiltonian.
+
+ Parameters
+ ----------
+ tb_hamiltonian : dict
+ Keys are tuples of integers: the offsets of the hoppings
+ ((0, 0, 0) for the onsite). Values are symbolic expressions
+ for the hoppings/onsite or expressions that can by sympified using
+ `kwant.continuum.sympify`.
+ discrete_coordinates : sequence of strings
+ Set of coordinates for which momentum operators will be treated as
+ differential operators. For example ``discrete_coordinates=('x', 'y')``.
+ lattice_constant : int or float, default: 1
+ Lattice constant for the template Builder.
+ substitutions : dict, defaults to empty
+ A namespace to be passed to `kwant.continuum.sympify` when
+ ``tb_hamiltonian`` is a string. Could be used to simplify matrix input:
+ ``sympify('k_x**2 * s_z', substitutions={'s_z': [[1, 0], [0, -1]]})``.
+ verbose : bool, default: False
+ If ``True`` additional information will be printed.
+
+ Returns
+ -------
+ `~kwant.builder.Builder` with translational symmetry,
+ which can be used as a template.
+ """
+ if len(discrete_coordinates) == 0:
+ raise ValueError('Discrete coordinates cannot be empty.')
+
+ for k, v in tb_hamiltonian.items():
+ if not isinstance(v, (sympy.Expr, sympy.matrices.MatrixBase)):
+ tb_hamiltonian[k] = sympify(v, substitutions)
+
+ discrete_coordinates = sorted(discrete_coordinates)
+
+ tb = {}
+ for n, (offset, hopping) in enumerate(tb_hamiltonian.items()):
+ if verbose:
+ print("Function generated for {}:".format(offset))
+
+ onsite = all(i == 0 for i in offset)
+
+ if onsite:
+ name = 'onsite'
+ else:
+ name = 'hopping_{}'.format(n)
+
+ tb[offset] = _value_function(hopping, discrete_coordinates,
+ lattice_constant, onsite, name,
+ verbose=verbose)
+
+ dim = len(discrete_coordinates)
+ onsite_zeros = (0,) * dim
+
+ prim_vecs = lattice_constant * np.eye(dim)
+ random_element = next(iter(tb_hamiltonian.values()))
+ norbs = (1 if isinstance(random_element, sympy.Expr)
+ else random_element.shape[0])
+ lattice = Monatomic(prim_vecs, norbs=norbs)
+
+ onsite = tb.pop(onsite_zeros)
+ # 'delta' parameter to HoppingKind is the negative of the 'hopping offset'
+ hoppings = {HoppingKind(tuple(-i for i in d), lattice): val
+ for d, val in tb.items()}
+
+ syst = Builder(TranslationalSymmetry(*prim_vecs))
+ syst[lattice(*onsite_zeros)] = onsite
+ for hop, val in hoppings.items():
+ syst[hop] = val
+
+ return syst
+
+
+def _differentiate(expression, coordinate_name):
+ """Calculate derivative of an expression for given coordinate.
+
+ Parameters:
+ -----------
+ expression : sympy.Expr
+ Sympy expression containing function to be derivated.
+ coordinate_name : string
+ Coordinate over which derivative is calculated.
+
+ Returns
+ -------
+ sympy.Expr
+ """
+ assert coordinate_name in 'xyz'
+ coordinate = _position_operators[coordinate_name]
+ h = _displacements[coordinate_name]
+
+ expr1 = expression.subs(coordinate, coordinate + h)
+ expr2 = expression.subs(coordinate, coordinate - h)
+
+ return (expr1 - expr2) / (2 * h)
+
+
+def _discretize_summand(summand, discrete_coordinates):
+ """Discretize a product of factors.
+
+ Parameters
+ ----------
+ summand : sympy.Expr
+ discrete_coordinates : sequence of strings
+ Must be a subset of ``{'x', 'y', 'z'}``.
+
+ Returns
+ -------
+ sympy.Expr
+ """
+ assert not isinstance(summand, sympy.Add), "Input should be one summand."
+ momenta = ['k_{}'.format(s) for s in discrete_coordinates]
+
+ factors = reversed(summand.as_ordered_factors())
+ result = 1
+ for factor in factors:
+ symbol, exponent = factor.as_base_exp()
+ if isinstance(symbol, sympy.Symbol) and symbol.name in momenta:
+ for i in range(exponent):
+ coordinate = symbol.name[-1]
+ # apply momentum as differential operator '-i d/dx'
+ result = -sympy.I * _differentiate(result, coordinate)
+ else:
+ result = factor * result
+
+ return sympy.expand(result)
+
+
+def _discretize_expression(expression, discrete_coordinates):
+ """Discretize an expression into a discrete (tight-binding) representation.
+
+ Parameters
+ ----------
+ expression : sympy.Expr
+ discrete_coordinates : sequence of strings
+ Must be a subset of ``{'x', 'y', 'z'}``.
+
+ Returns
+ -------
+ dict
+ Keys are tuples of integers; the offsets of the hoppings
+ ((0, 0, 0) for the onsite). Values are symbolic expressions
+ for the hoppings/onsite.
+ """
+ def _read_offset(wf):
+ # e.g. wf(x + h, y, z + h) -> (1, 0, 1)
+ assert wf.func == _wf
+
+ offset = []
+ for c, arg in zip(discrete_coordinates, wf.args):
+ coefficients = arg.as_coefficients_dict()
+ assert coefficients[_position_operators[c]] == 1
+
+ ai = _displacements[c]
+ offset.append(coefficients.pop(ai, 0))
+ return tuple(offset)
+
+ def _extract_hoppings(expr):
+ """Read hoppings and perform shortening operation."""
+ expr = sympy.expand(expr)
+ summands = expr.args if expr.func == sympy.Add else [expr]
+
+ offset = [_read_offset(s.args[-1]) for s in summands]
+ coeffs = [sympy.Mul(*s.args[:-1]) for s in summands]
+ offset = np.array(offset, dtype=int)
+ # rescale the offsets for each coordinate by their greatest
+ # common divisor across the summands. e.g:
+ # wf(x+2h) + wf(x+4h) --> wf(x+h) + wf(x+2h) and a_x //= 2
+ subs = {}
+ for i, xi in enumerate(discrete_coordinates):
+ factor = int(gcd(*offset[:, i]))
+ if factor < 1:
+ continue
+ offset[:, i] //= factor
+ subs[_displacements[xi]] = sympy.symbols('a') / factor
+ # apply the rescaling to the hoppings
+ output = defaultdict(lambda: sympy.Integer(0))
+ for n, c in enumerate(coeffs):
+ output[tuple(offset[n].tolist())] += c.subs(subs)
+ return dict(output)
+
+ # main function body starts here
+ if (isinstance(expression, (int, float, sympy.Integer, sympy.Float))
+ or not expression.atoms(sympy.Symbol) & set(momentum_operators)):
+ n = len(discrete_coordinates)
+ return {(0,) * n: expression}
+
+ assert isinstance(expression, sympy.Expr)
+
+ # make sure we have list of summands
+ summands = expression.as_ordered_terms()
+
+ # discretize every summand
+ coordinates = tuple(_position_operators[s] for s in discrete_coordinates)
+ wf = _wf(*coordinates)
+
+ discrete_expression = defaultdict(int)
+ for summand in summands:
+ summand = _discretize_summand(summand * wf, discrete_coordinates)
+ hops = _extract_hoppings(summand)
+ for k, v in hops.items():
+ discrete_expression[k] += v
+
+ return dict(discrete_expression)
+
+
+################ string processing
+
+class _NumericPrinter(LambdaPrinter):
+
+ def _print_ImaginaryUnit(self, expr):
+ # prevent sympy from printing 'I' for imaginary unit
+ return "1j"
+
+ def _print_Pow(self, expr):
+ # copied from sympy's StrPrinter with the code paths
+ # to print 'sqrt' removed.
+ PREC = precedence(expr)
+
+ if expr.is_commutative and expr.exp is -sympy.S.One:
+ return '1/%s' % self.parenthesize(expr.base, PREC)
+
+ e = self.parenthesize(expr.exp, PREC)
+ if (self.printmethod == '_sympyrepr' and
+ expr.exp.is_Rational and expr.exp.q != 1):
+ if e.startswith('(Rational'):
+ return '%s**%s' % (self.parenthesize(expr.base, PREC), e[1:-1])
+ return '%s**%s' % (self.parenthesize(expr.base, PREC), e)
+
+
+def _print_sympy(expr):
+ return lambdastr((), expr, printer=_NumericPrinter)[len('lambda : '):]
+
+
+def _return_string(expr, discrete_coordinates):
+ """Process a sympy expression into an evaluatable Python return statement.
+
+ Parameters
+ ----------
+ expr : sympy.Expr
+
+ Returns
+ -------
+ output : string
+ A return string that can be used to assemble a Kwant value function.
+ map_function_calls : dict
+ mapping of function calls to assigned constants.
+ const_symbols : sequance of sympy.Symbol
+ All constants that appear in the expression.
+ _cache: dict
+ mapping of cache symbols to cached matrices.
+ """
+ _cache = {}
+ def cache(x):
+ s = sympy.symbols('_cache_{}'.format(len(_cache)))
+ _cache[str(s)] = ta.array(x.tolist(), complex)
+ return s
+
+ blacklisted = set(discrete_coordinates) | {'site', 'site1', 'site2'}
+ const_symbols = {s for s in expr.atoms(sympy.Symbol)
+ if s.name not in blacklisted}
+
+ # functions will be evaluated within the function body and the
+ # result assigned to a symbol '_const_<n>', so we replace all
+ # function calls by these symbols in the return statement.
+ map_function_calls = expr.atoms(AppliedUndef, sympy.Function)
+ map_function_calls = {s: sympy.symbols('_const_{}'.format(n))
+ for n, s in enumerate(map_function_calls)}
+
+ expr = expr.subs(map_function_calls)
+
+ if isinstance(expr, sympy.matrices.MatrixBase):
+ # express matrix return values in terms of sums of known matrices,
+ # which will be assigned to '_cache_n' in the function body.
+ mons = matrix_monomials(expr, *expr.atoms(sympy.Symbol))
+ mons = {k: cache(v) for k, v in mons.items()}
+ mons = ["{} * {}".format(_print_sympy(k), _print_sympy(v))
+ for k, v in mons.items()]
+ output = " + ".join(mons)
+ else:
+ output = _print_sympy(expr)
+
+ return 'return {}'.format(output), map_function_calls, const_symbols, _cache
+
+
+def _assign_symbols(map_function_calls, lattice_constant,
+ discrete_coordinates, onsite):
+ """Generate a series of assignments.
+
+ Parameters
+ ----------
+ map_function_calls : dict
+ mapping of function calls to assigned constants.
+ lattice_constant : int or float
+ Used to get site.pos from site.tag
+ discrete_coordinates : sequence of strings
+ If left as None coordinates will not be read from a site.
+ onsite : bool
+ True if function is called for onsite, false for hoppings
+
+ Returns
+ -------
+ assignments : list of strings
+ List of lines used for including in a function.
+ """
+ lines = []
+
+ if discrete_coordinates:
+ site = 'site' if onsite else 'site1'
+ args = ', '.join(discrete_coordinates), str(lattice_constant), site
+ lines.append('({}, ) = {} * {}.tag'.format(*args))
+
+ for k, v in map_function_calls.items():
+ lines.append("{} = {}".format(v, _print_sympy(k)))
+
+ return lines
+
+
+def _value_function(expr, discrete_coordinates, lattice_constant, onsite,
+ name='_anonymous_func', verbose=False):
+ """Generate a numeric function from a sympy expression.
+
+ Parameters
+ ----------
+ expr : sympy.Expr or sympy.matrix
+ Expr that from which value function will be generated.
+ discrete_coordinates : sequence of strings
+ List of coodinates present in the system.
+ lattice_constant : int or float
+ Lattice spacing of the system
+ verbose : bool, default: False
+ If True, the function body is printed.
+
+ Returns
+ -------
+ numerical function that can be used with Kwant.
+ """
+
+ expr = expr.subs({sympy.Symbol('a'): lattice_constant})
+ return_string, map_function_calls, const_symbols, _cache = \
+ _return_string(expr, discrete_coordinates=discrete_coordinates)
+
+ # first check if value function needs to read coordinates
+ if not ({_position_operators[s] for s in discrete_coordinates} &
+ expr.atoms(sympy.Symbol)):
+ discrete_coordinates = None
+
+ # constants and functions in the sympy input will be passed
+ # as keyword-only arguments to the value function
+ required_kwargs = set.union({s.name for s in const_symbols},
+ {str(k.func) for k in map_function_calls})
+ required_kwargs = ', '.join(sorted(required_kwargs))
+
+ if (not required_kwargs) and (discrete_coordinates is None):
+ # we can just use a constant value instead of a value function
+ if isinstance(expr, sympy.MatrixBase):
+ output = ta.array(expr.tolist(), complex)
+ else:
+ output = complex(expr)
+
+ if verbose:
+ print("\n{}\n\n".format(output))
+
+ return output
+
+ lines = _assign_symbols(map_function_calls, onsite=onsite,
+ lattice_constant=lattice_constant,
+ discrete_coordinates=discrete_coordinates)
+
+ lines.append(return_string)
+
+ separator = '\n' + 4 * ' '
+ # 'site_string' is tightly coupled to the symbols used in '_assign_symbol'
+ site_string = 'site' if onsite else 'site1, site2'
+ if required_kwargs:
+ header_str = 'def {}({}, *, {}):'
+ header = header_str.format(name, site_string, required_kwargs)
+ else:
+ header = 'def {}({}):'.format(name, site_string)
+ func_code = separator.join([header] + list(lines))
+
+ namespace = {'pi': np.pi}
+ namespace.update(_cache)
+
+ if verbose:
+ for k, v in _cache.items():
+ print("\n{} = (\n{})".format(k, repr(np.array(v))))
+ print('\n' + func_code + '\n\n')
+
+ exec(func_code, namespace)
+ f = namespace[name]
+
+ return f
diff --git a/kwant/continuum/tests/__init__.py b/kwant/continuum/tests/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/kwant/continuum/tests/test_common.py b/kwant/continuum/tests/test_common.py
new file mode 100644
index 0000000000000000000000000000000000000000..fa562c1679bfbdbde09e00f7dc916805b7f513c3
--- /dev/null
+++ b/kwant/continuum/tests/test_common.py
@@ -0,0 +1,146 @@
+from kwant.continuum._common import position_operators, momentum_operators
+from kwant.continuum._common import make_commutative, sympify
+from kwant.continuum._common import expression_monomials, matrix_monomials
+from kwant.continuum._common import lambdify
+
+import tinyarray as ta
+from sympy.physics.matrices import msigma
+import sympy
+from functools import reduce
+from operator import mul
+import pytest
+
+
+def test_sympify():
+ A, B, C = sympy.symbols('A B C')
+ x, y, z = position_operators
+ kx, ky, kz = momentum_operators
+
+ # basics
+ assert sympify('k_x * A(x) * k_x') == kx * A(x) * kx
+ assert sympify('[[k_x * A(x) * k_x]]') == sympy.Matrix([kx * A(x) * kx])
+
+ # using substitutions
+ symbolic_pauli = {'sigma_x': msigma(1), 'sigma_y': msigma(2), 'sigma_z': msigma(3)}
+ got = sympify('k_x * sigma_y + k_y * sigma_x', substitutions=symbolic_pauli)
+ assert got == kx * symbolic_pauli['sigma_y'] + ky * symbolic_pauli['sigma_x']
+
+ got = sympify("sigma_y", substitutions={'sigma_y': "[[0, -1j], [1j, 0]]"})
+ assert got == symbolic_pauli['sigma_y']
+
+ got = sympify("sigma_y", substitutions={'sigma_y': [[0, -sympy.I], [sympy.I, 0]]})
+ assert got == symbolic_pauli['sigma_y']
+
+ got = sympify('[[k_x*A(x)*k_x, B(x, y)*k_x], [k_x*B(x, y), C*k_y**2]]')
+ assert got == sympy.Matrix([[kx*A(x)*kx, B(x, y)*kx], [kx*B(x, y), C*ky**2]])
+
+
+
+
+A, B, non_x = sympy.symbols('A B x', commutative=False)
+x = sympy.Symbol('x')
+
+expr1 = non_x*A*non_x + x**2 * A * x + B*non_x**2
+
+matr = sympy.Matrix([[expr1, expr1+A*non_x], [0, -expr1]])
+res_mat = sympy.Matrix([[x**3*A + x**2*A + x**2*B, x**3*A + x**2*A + x**2*B + x*A],
+ [0, -x**3*A - x**2*A - x**2*B]])
+
+def test_make_commutative():
+ assert make_commutative(expr1, x) == make_commutative(expr1, non_x)
+ assert make_commutative(expr1, x) == x**3*A + x**2*A + x**2*B
+ assert make_commutative(matr, x) == res_mat
+
+
+expr2 = non_x*A*non_x + x**2 * A*2 * x + B*non_x/2 + non_x*B/2 + x + A + non_x + x/A
+
+def test_expression_monomials():
+ assert expression_monomials(expr2, x) == {x**3: 2*A, 1: A, x: 2 + A**(-1) + B, x**2: A}
+ assert expression_monomials(expr1, x) == {x**2: A + B, x**3: A}
+ assert expression_monomials(x, x) == {x: 1}
+ assert expression_monomials(x**2, x) == {x**2: 1}
+ assert expression_monomials(x**2 + x, x) == {x: 1, x**2: 1}
+
+ expr = 1 + x + A*x + 2*x + x**2 + A*x**2 + non_x*A*non_x
+ out = {1: 1, x: 3 + A, x**2: 2 * A + 1}
+ assert expression_monomials(expr, x) == out
+
+ expr = 1 + x * (3 + A) + x**2 * (1 + A)
+ out = {1: 1, x: 3 + A, x**2: 1 * A + 1}
+ assert expression_monomials(expr, x) == out
+
+
+def legacy_expression_monomials(expr, *gens):
+ """ This was my first implementation. Unfortunately it is very slow.
+
+ It is used to test correctness of new matrix_monomials function.
+ """
+ expr = make_commutative(expr, x)
+ R = sympy.ring(gens, sympy.EX, sympy.lex)[0]
+ expr = R(expr)
+
+ output = {}
+ for power, coeff in zip(expr.monoms(), expr.coeffs()):
+ key = reduce(mul, [sympy.Symbol(k.name)**n for k, n in zip(gens, power)])
+ output[key] = sympy.expand(coeff.as_expr())
+ return output
+
+
+def test_expression_monomials_with_reference_function():
+ assert legacy_expression_monomials(expr2, x) == expression_monomials(expr2, x)
+
+
+
+def test_matrix_monomials():
+ out = {
+ x**2: sympy.Matrix([[A + B, A + B],[0, -A - B]]),
+ x: sympy.Matrix([[0, A], [0, 0]]),
+ x**3: sympy.Matrix([[A, A], [0, -A]])}
+ mons = matrix_monomials(matr, x)
+ assert mons == out
+
+
+
+@pytest.mark.parametrize("e, should_be, kwargs", [
+ ("x+y", lambda x, y: x+y, dict(x=1, y=2)),
+ ("1", lambda: 1, dict()),
+ ("f(x)", lambda f, x: f(x), dict(f=lambda x: x, x=2)),
+ (sympify("f(x)"), lambda f, x: f(x), dict(f=lambda x: x, x=2)),
+ ("[[f(x)]]", lambda f, x: ta.array(f(x)), dict(f=lambda x: x, x=2))
+])
+def test_lambdify(e, should_be, kwargs):
+ e = lambdify(e)
+ assert e(**kwargs) == should_be(**kwargs)
+
+
+
+
+
+# dispersion_string = ('A * k_x**2 * eye(2) + B * k_y**2 * eye(2)'
+# '+ alpha * k_x * sigma_y')
+
+# def dispersion_function(kx, ky, A, B, alpha, **kwargs):
+# h_k = ((A * kx**2 + B * ky**2) * sigma_0 +
+# alpha * kx * sigma_y)
+# return np.linalg.eigvalsh(h_k)
+
+# dispersion_kwargs = {'A': 1, 'B': 2, 'alpha': 0.5, 'M': 2}
+# sigma_0 = np.eye(2)
+# sigma_y = np.array([[0, -1j], [1j, 0]])
+
+
+# @pytest.mark.parametrize("expr, should_be, kwargs", [
+# (dispersion_string, dispersion_function, dispersion_kwargs),
+# (sympify(dispersion_string), dispersion_function, dispersion_kwargs),
+# ])
+# def test_lambdify(expr, should_be, kwargs):
+# N = 5
+
+# x = y = np.linspace(-N, N+1)
+# xy = list(itertools.product(x, y))
+
+# h_k = lambdify(expr)
+# energies = [la.eigvalsh(h_k(k_x=kx, k_y=ky, **kwargs))
+# for kx, ky in xy]
+# energies_should_be = [should_be(kx, ky, **kwargs) for kx, ky in xy]
+# assert np.allclose(energies, energies_should_be)
diff --git a/kwant/continuum/tests/test_discretizer.py b/kwant/continuum/tests/test_discretizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..014c37b32f606057b7990dad58787b4ee6096a39
--- /dev/null
+++ b/kwant/continuum/tests/test_discretizer.py
@@ -0,0 +1,480 @@
+import sympy
+import numpy as np
+
+from ..discretizer import discretize
+from ..discretizer import discretize_symbolic
+from ..discretizer import build_discretized
+from ..discretizer import _wf
+
+import inspect
+from functools import wraps
+
+
+def swallows_extra_kwargs(f):
+ sig = inspect.signature(f)
+ pars = sig.parameters
+ if any(i.kind is inspect.Parameter.VAR_KEYWORD for i in pars.values()):
+ return f
+
+ names = {name for name, value in pars.items() if
+ value.kind not in (inspect.Parameter.VAR_POSITIONAL,
+ inspect.Parameter.POSITIONAL_ONLY)}
+
+ @wraps(f)
+ def wrapped(*args, **kwargs):
+ bound = sig.bind(*args, **{name: value for name, value
+ in kwargs.items() if name in names})
+ return f(*bound.args, **bound.kwargs)
+
+ return wrapped
+
+
+I = sympy.I
+kx, ky, kz = sympy.symbols('k_x k_y k_z', commutative=False)
+x, y, z = sympy.symbols('x y z', commutative=False)
+ax, ay, az = sympy.symbols('a_x a_y a_z')
+a = sympy.symbols('a')
+
+wf = _wf
+Psi = wf(x, y, z)
+A, B = sympy.symbols('A B', commutative=False)
+
+ns = {'A': A, 'B': B, 'a_x': ax, 'a_y': ay, 'az': az, 'x': x, 'y': y, 'z': z}
+
+
+def test_reading_coordinates():
+ test = {
+ kx**2 : {'x'},
+ kx**2 + ky**2 : {'x', 'y'},
+ kx**2 + ky**2 + kz**2 : {'x', 'y', 'z'},
+ ky**2 + kz**2 : {'y', 'z'},
+ kz**2 : {'z'},
+ kx * A(x,y) * kx : {'x'},
+ kx**2 + kz * B(y) : {'x', 'z'},
+ }
+ for inp, out in test.items():
+ ham, got = discretize_symbolic(inp)
+ assert all(d in out for d in got),\
+ "Should be: _split_factors({})=={}. Not {}".format(inp, out, got)
+
+
+def test_reading_coordinates_matrix():
+ test = [
+ (sympy.Matrix([kx**2]) , {'x'}),
+ (sympy.Matrix([kx**2 + ky**2]) , {'x', 'y'}),
+ (sympy.Matrix([kx**2 + ky**2 + kz**2]) , {'x', 'y', 'z'}),
+ (sympy.Matrix([ky**2 + kz**2]) , {'y', 'z'}),
+ (sympy.Matrix([kz**2]) , {'z'}),
+ (sympy.Matrix([kx * A(x,y) * kx]) , {'x'}),
+ (sympy.Matrix([kx**2 + kz * B(y)]) , {'x', 'z'}),
+ ]
+ for inp, out in test:
+ ham, got = discretize_symbolic(inp)
+ assert all(d in out for d in got),\
+ "Should be: _split_factors({})=={}. Not {}".format(inp, out, got)
+
+
+def test_reading_different_matrix_types():
+ test = [
+ (sympy.MutableMatrix([kx**2]) , {'x'}),
+ (sympy.ImmutableMatrix([kx**2]) , {'x'}),
+ (sympy.MutableDenseMatrix([kx**2]) , {'x'}),
+ (sympy.ImmutableDenseMatrix([kx**2]) , {'x'}),
+ ]
+ for inp, out in test:
+ ham, got = discretize_symbolic(inp)
+ assert all(d in out for d in got),\
+ "Should be: _split_factors({})=={}. Not {}".format(inp, out, got)
+
+
+def test_simple_derivations():
+ test = {
+ kx**2 : {(0,): 2/a**2, (1,): -1/a**2},
+ kx**2 + ky**2 : {(0, 1): -1/a**2, (0, 0): 4/a**2,
+ (1, 0): -1/a**2},
+ kx**2 + ky**2 + kz**2 : {(1, 0, 0): -1/a**2, (0, 0, 1): -1/a**2,
+ (0, 0, 0): 6/a**2, (0, 1, 0): -1/a**2},
+ ky**2 + kz**2 : {(0, 1): -1/a**2, (0, 0): 4/a**2,
+ (1, 0): -1/a**2},
+ kz**2 : {(0,): 2/a**2, (1,): -1/a**2},
+ kx * A(x,y) * kx : {(1, ): -A(a/2 + x, y)/a**2,
+ (0, ): A(-a/2 + x, y)/a**2 + A(a/2 + x, y)/a**2},
+ kx**2 + kz * B(y) : {(1, 0): -1/a**2, (0, 1): -I*B(y)/(2*a),
+ (0, 0): 2/a**2},
+ kx * A(x) : {(0,): 0, (1,): -I*A(a + x)/(2*a)},
+ ky * A(x) : {(1,): -I*A(x)/(2*a), (0,): 0},
+ kx * A(x) * B : {(0,): 0, (1,): -I*A(a + x)*B/(2*a)},
+ 5 * kx : {(0,): 0, (1,): -5*I/(2*a)},
+ kx * (A(x) + B(x)) : {(0,): 0,
+ (1,): -I*A(a + x)/(2*a) - I*B(a + x)/(2*a)},
+ }
+ for inp, out in test.items():
+ got, _ = discretize_symbolic(inp)
+ assert got == out
+
+ for inp, out in test.items():
+ got, _ = discretize_symbolic(str(inp), substitutions=ns)
+ assert got == out
+
+
+def test_simple_derivations_matrix():
+ test = {
+ kx**2 : {(0,): 2/a**2, (1,): -1/a**2},
+ kx**2 + ky**2 : {(0, 1): -1/a**2, (0, 0): 4/a**2,
+ (1, 0): -1/a**2},
+ kx**2 + ky**2 + kz**2 : {(1, 0, 0): -1/a**2, (0, 0, 1): -1/a**2,
+ (0, 0, 0): 6/a**2, (0, 1, 0): -1/a**2},
+ ky**2 + kz**2 : {(0, 1): -1/a**2, (0, 0): 4/a**2,
+ (1, 0): -1/a**2},
+ kz**2 : {(0,): 2/a**2, (1,): -1/a**2},
+ kx * A(x,y) * kx : {(1, ): -A(a/2 + x, y)/a**2,
+ (0, ): A(-a/2 + x, y)/a**2 + A(a/2 + x, y)/a**2},
+ kx**2 + kz * B(y) : {(1, 0): -1/a**2, (0, 1): -I*B(y)/(2*a),
+ (0, 0): 2/a**2},
+ kx * A(x) : {(0,): 0, (1,): -I*A(a + x)/(2*a)},
+ ky * A(x) : {(1,): -I*A(x)/(2*a), (0,): 0},
+ kx * A(x) * B : {(0,): 0, (1,): -I*A(a + x)*B/(2*a)},
+ 5 * kx : {(0,): 0, (1,): -5*I/(2*a)},
+ kx * (A(x) + B(x)) : {(0,): 0,
+ (1,): -I*A(a + x)/(2*a) - I*B(a + x)/(2*a)},
+ }
+
+ new_test = []
+ for inp, out in test.items():
+ new_out = {}
+ for k, v in out.items():
+ new_out[k] = sympy.Matrix([v])
+ new_test.append((sympy.Matrix([inp]), new_out))
+
+ for inp, out in new_test:
+ got, _ = discretize_symbolic(inp)
+ assert got == out
+
+ for inp, out in new_test:
+ got, _ = discretize_symbolic(str(inp), substitutions=ns)
+ assert got == out
+
+ for inp, out in new_test:
+ got, _ = discretize_symbolic(str(inp).replace('Matrix', ''), substitutions=ns)
+ assert got == out
+
+
+
+def test_integer_float_input():
+ test = {
+ 0 : {(0,0,0): 0},
+ 1 : {(0,0,0): 1},
+ 5 : {(0,0,0): 5},
+ }
+
+ for inp, out in test.items():
+ got, _ = discretize_symbolic(int(inp), {'x', 'y', 'z'})
+ assert got == out
+
+ got, _ = discretize_symbolic(float(inp), {'x', 'y', 'z'})
+ assert got == out
+
+ # let's test in matrix version too
+ new_test = []
+ for inp, out in test.items():
+ new_out = {}
+ for k, v in out.items():
+ new_out[k] = sympy.Matrix([v])
+ new_test.append((inp, new_out))
+
+ for inp, out in new_test:
+ got, _ = discretize_symbolic(sympy.Matrix([int(inp)]), {'x', 'y', 'z'})
+ assert got == out
+
+ got, _ = discretize_symbolic(sympy.Matrix([float(inp)]), {'x', 'y', 'z'})
+ assert got == out
+
+
+def test_different_discrete_coordinates():
+ test = [
+ (
+ {'x', 'y', 'z'}, {
+ (1, 0, 0): -1/a**2, (0, 0, 1): -1/a**2,
+ (0, 0, 0): 6/a**2, (0, 1, 0): -1/a**2
+ }
+ ),
+ (
+ {'x', 'y'}, {
+ (0, 1): -1/a**2,
+ (1, 0): -1/a**2,
+ (0, 0): kz**2 + 4/a**2
+ }
+ ),
+ (
+ {'x', 'z'}, {
+ (0, 1): -1/a**2,
+ (1, 0): -1/a**2,
+ (0, 0): ky**2 + 4/a**2
+ }
+ ),
+ (
+ {'y', 'z'}, {
+ (0, 1): -1/a**2,
+ (1, 0): -1/a**2,
+ (0, 0): kx**2 + 4/a**2
+ }
+ ),
+ (
+ {'x'}, {
+ (0,): ky**2 + kz**2 + 2/a**2, (1,): -1/a**2
+ }
+ ),
+ (
+ {'y'}, {
+ (0,): kx**2 + kz**2 + 2/a**2, (1,): -1/a**2
+ }
+ ),
+ (
+ {'z'}, {
+ (0,): ky**2 + kx**2 + 2/a**2, (1,): -1/a**2
+ }
+ ) ,
+ ]
+ for inp, out in test:
+ got, _ = discretize_symbolic(kx**2 + ky**2 + kz**2, inp)
+ assert got == out
+
+ # let's test in matrix version too
+ new_test = []
+ for inp, out in test:
+ new_out = {}
+ for k, v in out.items():
+ new_out[k] = sympy.Matrix([v])
+ new_test.append((inp, new_out))
+
+ for inp, out in new_test:
+ got, _ = discretize_symbolic(sympy.Matrix([kx**2 + ky**2 + kz**2]), inp)
+ assert got == out
+
+
+def test_non_expended_input():
+ symbolic, coords = discretize_symbolic(kx * (kx + A(x)))
+ desired = {
+ (0,): 2/a**2,
+ (1,): -I*A(a + x)/(2*a) - 1/a**2
+ }
+ assert symbolic == desired
+
+
+
+def test_matrix_with_zeros():
+ Matrix = sympy.Matrix
+ symbolic, _ = discretize_symbolic("[[k_x*A(x)*k_x, 0], [0, k_x*A(x)*k_x]]")
+ output = {
+ (0,) : Matrix([[A(-a/2 + x)/a**2 + A(a/2 + x)/a**2, 0], [0, A(-a/2 + x)/a**2 + A(a/2 + x)/a**2]]),
+ (1,) : Matrix([[-A(a/2 + x)/a**2, 0], [0, -A(a/2 + x)/a**2]]),
+ }
+ assert symbolic == output
+
+
+def test_numeric_functions_basic_symbolic():
+ for i in [0, 1, 3, 5]:
+ builder = discretize(i, {'x'})
+ lat = next(iter(builder.sites()))[0]
+ assert builder[lat(0)] == i
+
+ p = dict(t=i)
+
+ tb = {(0,): sympy.sympify("2*t"), (1,): sympy.sympify('-t')}
+ builder = build_discretized(tb, {'x'}, lattice_constant=1)
+ lat = next(iter(builder.sites()))[0]
+ assert 2*p['t'] == builder[lat(0)](None, **p)
+ assert -p['t'] == builder[lat(1), lat(0)](None, None, **p)
+
+ tb = {(0,): sympy.sympify("0"), (1,): sympy.sympify('-1j * t')}
+ builder = build_discretized(tb, {'x'}, lattice_constant=1)
+ lat = next(iter(builder.sites()))[0]
+ assert -1j * p['t'] == builder[lat(0), lat(1)](None, None, **p)
+ assert +1j * p['t'] == builder[lat(1), lat(0)](None, None, **p)
+
+
+def test_numeric_functions_not_discrete_coords():
+ builder = discretize('k_y + y', 'x')
+ lat = next(iter(builder.sites()))[0]
+ onsite = builder[lat(0)]
+
+ assert onsite(None, k_y=2, y=1) == 2 + 1
+
+
+def test_numeric_functions_with_pi():
+ # Two cases because once it is casted
+ # to complex, one there is a function created
+
+ builder = discretize('A + pi', 'x')
+ lat = next(iter(builder.sites()))[0]
+ onsite = builder[lat(0)]
+ assert onsite(None, A=1) == 1 + np.pi
+
+
+ builder = discretize('pi', 'x')
+ lat = next(iter(builder.sites()))[0]
+ onsite = builder[lat(0)]
+ assert onsite == np.pi
+
+
+def test_numeric_functions_basic_string():
+ for i in [0, 1, 3, 5]:
+ builder = discretize(i, {'x'})
+ lat = next(iter(builder.sites()))[0]
+ assert builder[lat(0)] == i
+
+ p = dict(t=i)
+
+ tb = {(0,): "2*t", (1,): "-t"}
+ builder = build_discretized(tb, {'x'}, lattice_constant=1)
+ lat = next(iter(builder.sites()))[0]
+ assert 2*p['t'] == builder[lat(0)](None, **p)
+ assert -p['t'] == builder[lat(1), lat(0)](None, None, **p)
+
+ tb = {(0,): "0", (1,): "-1j * t"}
+ builder = build_discretized(tb, {'x'}, lattice_constant=1)
+ lat = next(iter(builder.sites()))[0]
+ assert -1j * p['t'] == builder[lat(0), lat(1)](None, None, **p)
+ assert +1j * p['t'] == builder[lat(1), lat(0)](None, None, **p)
+
+ tb = {(0,): "0", (-1,): "+1j * t"}
+ builder = build_discretized(tb, {'x'}, lattice_constant=1)
+ lat = next(iter(builder.sites()))[0]
+ assert -1j * p['t'] == builder[lat(0), lat(1)](None, None, **p)
+ assert +1j * p['t'] == builder[lat(1), lat(0)](None, None, **p)
+
+
+def test_numeric_functions_advance():
+ hams = [
+ kx**2,
+ kx**2 + x,
+ A(x),
+ kx*A(x)*kx,
+ sympy.Matrix([[kx * A(x) * kx, A(x)*kx], [kx*A(x), A(x)+B]]),
+ kx**2 + B * x,
+ 'k_x**2 + sin(x)',
+ B ** 0.5 * kx**2,
+ B ** (1/2) * kx**2,
+ sympy.sqrt(B) * kx**2,
+
+ ]
+ for hamiltonian in hams:
+ for a in [1, 2, 5]:
+ for fA in [lambda x: x, lambda x: x**2, lambda x: x**3]:
+ symbolic, coords = discretize_symbolic(hamiltonian, {'x'})
+ builder = build_discretized(symbolic, coords, lattice_constant=a)
+ lat = next(iter(builder.sites()))[0]
+
+ p = dict(A=fA, B=5, sin=np.sin)
+
+ # test onsite
+ v = symbolic.pop((0,)).subs({sympy.symbols('a'): a, B: p['B']})
+ f_sym = sympy.lambdify(['A', 'x'], v)
+ f_num = builder[lat(0)]
+
+ if callable(f_num):
+ f_num = swallows_extra_kwargs(f_num)
+ for n in range(-100, 100, 10):
+ assert np.allclose(f_sym(fA, a*n), f_num(lat(n), **p))
+ else:
+ for n in range(-100, 100, 10):
+ assert np.allclose(f_sym(fA, a*n), f_num)
+
+
+ # test hoppings
+ for k, v in symbolic.items():
+ v = v.subs({sympy.symbols('a'): a, B: p['B']})
+ f_sym = sympy.lambdify(['A', 'x'], v)
+ f_num = builder[lat(0), lat(k[0])]
+
+ if callable(f_num):
+ f_num = swallows_extra_kwargs(f_num)
+ for n in range(10):
+ lhs = f_sym(fA, a * n)
+ rhs = f_num(lat(n), lat(n+k[0]), **p)
+ assert np.allclose(lhs, rhs)
+ else:
+ for n in range(10):
+ lhs = f_sym(fA, a * n)
+ rhs = f_num
+ assert np.allclose(lhs, rhs)
+
+
+def test_numeric_functions_with_parameter():
+ hams = [
+ kx**2 + A(B, x)
+ ]
+ for hamiltonian in hams:
+ for a in [1, 2, 5]:
+ for fA in [lambda c, x: x+c, lambda c, x: x**2 + c]:
+ symbolic, coords = discretize_symbolic(hamiltonian, {'x'})
+ builder = build_discretized(symbolic, coords, lattice_constant=a)
+ lat = next(iter(builder.sites()))[0]
+
+ p = dict(A=fA, B=5)
+
+ # test onsite
+ v = symbolic.pop((0,)).subs({sympy.symbols('a'): a, B: p['B']})
+ f_sym = sympy.lambdify(['A', 'x'], v)
+
+ f_num = builder[lat(0)]
+ if callable(f_num):
+ f_num = swallows_extra_kwargs(f_num)
+
+
+ for n in range(10):
+ s = lat(n)
+ xi = a * n
+ if callable(f_num):
+ assert np.allclose(f_sym(fA, xi), f_num(s, **p))
+ else:
+ assert np.allclose(f_sym(fA, xi), f_num)
+
+
+ # test hoppings
+ for k, v in symbolic.items():
+ v = v.subs({sympy.symbols('a'): a, B: p['B']})
+ f_sym = sympy.lambdify(['A', 'x'], v)
+ f_num = builder[lat(0), lat(k[0])]
+
+ if callable(f_num):
+ f_num = swallows_extra_kwargs(f_num)
+
+ for n in range(10):
+ s = lat(n)
+ xi = a * n
+
+ lhs = f_sym(fA, xi)
+ if callable(f_num):
+ rhs = f_num(lat(n), lat(n+k[0]), **p)
+ else:
+ rhs = f_num
+
+ assert np.allclose(lhs, rhs)
+
+
+def test_basic_verbose(capsys): # or use "capfd" for fd-level
+ discretize('k_x * A(x) * k_x', verbose=True)
+ out, err = capsys.readouterr()
+ assert "Discrete coordinates set to" in out
+ assert "Function generated for (0,)" in out
+
+
+def test_that_verbose_covers_all_hoppings(capsys):
+ discretize('k_x**2 + k_y**2 + k_x*k_y', verbose=True)
+ out, err = capsys.readouterr()
+
+ for tag in [(0, 1), (0, 0), (1, -1), (1, 1)]:
+ assert "Function generated for {}".format(tag) in out
+
+
+def test_verbose_cache(capsys):
+ discretize('[[k_x * A(x) * k_x]]', verbose=True)
+ out, err = capsys.readouterr()
+ assert '_cache_0' in out
+
+
+def test_no_output_when_verbose_false(capsys):
+ discretize('[[k_x * A(x) * k_x]]', verbose=False)
+ out, err = capsys.readouterr()
+ assert out == ''