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/discretizer/__init__.py b/discretizer/__init__.py
deleted file mode 100644
index 441fa0c75e4af5d2ad29e0f5fb5cf60f8f05d8a3..0000000000000000000000000000000000000000
--- a/discretizer/__init__.py
+++ /dev/null
@@ -1,10 +0,0 @@
-__all__ = ['algorithms']
-
-
-from.discretizer import Discretizer
-import sympy
-
-__version__ = '0.4.3'
-
-momentum_operators = sympy.symbols('k_x k_y k_z', commutative=False)
-coordinates = sympy.symbols('x y z', commutative=False)
diff --git a/discretizer/algorithms.py b/discretizer/algorithms.py
deleted file mode 100644
index 6ad0e047092edefc8803b2e1b789f5f2954e0f38..0000000000000000000000000000000000000000
--- a/discretizer/algorithms.py
+++ /dev/null
@@ -1,416 +0,0 @@
-from __future__ import print_function, division
-
-import itertools
-import sympy
-import numpy as np
-from collections import defaultdict
-
-from .postprocessing import make_kwant_functions
-from .postprocessing import offset_to_direction
-
-from .interpolation import interpolate_tb_hamiltonian
-
-try:
- # normal situation
- from kwant.lattice import Monatomic
-except ImportError:
- # probably run on gitlab-ci
- pass
-
-# ************************** Some globals *********************************
-wavefunction_name = 'Psi'
-
-
-# **************** Operation on sympy expressions **************************
-def read_coordinates(expression):
- """Read coordinates used in expression.
-
- This function is used if ``discrete_coordinates`` are not provided by user.
-
- Parameters:
- -----------
- expression : sympy.Expr or sympy.Matrix instance
-
- Returns:
- --------
- discrete_coordinates : set of strings
- """
- discrete_coordinates = set()
- for a in expression.atoms(sympy.Symbol):
- if a.name in ['k_x', 'k_y', 'k_z']:
- discrete_coordinates.add(a.name.split('_')[1])
-
- if a.name in ['x', 'y', 'z']:
- discrete_coordinates.add(a.name)
-
- return discrete_coordinates
-
-def split_factors(expression, discrete_coordinates):
- """ Split symbolic `expression` for a discretization step.
-
- Parameters:
- -----------
- expression : sympy.Expr instance
- The expression to be split. It should represents single summand.
-
- Output:
- -------
- lhs : sympy.Expr instance
- Part of expression standing to the left from operators
- that acts in current discretization step.
-
- operators: sympy.Expr instance
- Operator that perform discretization in current step.
-
- rhs : sympy.Expr instance
- Part of expression that is derivated in current step.
-
- Raises:
- -------
- AssertionError
- if input `expression` is of type ``sympy.Add``
- """
- assert not isinstance(expression, sympy.Add), \
- 'Input expression must not be sympy.Add. It should be a single summand.'
-
- momentum_names = ['k_{}'.format(s) for s in discrete_coordinates]
- momentum_operators = sympy.symbols(momentum_names, commutative=False)
- momentum_operators += sympy.symbols(momentum_names, commutative=True)
-
- output = {'rhs': [1], 'operator': [1], 'lhs': [1]}
-
- if isinstance(expression, sympy.Pow):
- base, exponent = expression.args
- if base in momentum_operators:
- output['operator'].append(base)
- output['lhs'].append(sympy.Pow(base, exponent-1))
- else:
- output['rhs'].append(expression)
-
-
- elif isinstance(expression, (int, float, sympy.Integer, sympy.Float)):
- output['rhs'].append(expression)
-
- elif isinstance(expression, (sympy.Symbol, sympy.Function)):
- if expression in momentum_operators:
- output['operator'].append(expression)
- else:
- output['rhs'].append(expression)
-
- elif isinstance(expression, sympy.Mul):
- iterator = iter(expression.args[::-1])
- for factor in iterator:
- if factor in momentum_operators:
- output['operator'].append(factor)
- break
- elif factor.func == sympy.Pow and factor.args[0] in momentum_operators:
- base, exponent = factor.args
- output['operator'].append(base)
- output['lhs'].append(sympy.Pow(base, exponent-1))
- break
- else:
- output['rhs'].append(factor)
-
- for factor in iterator:
- output['lhs'].append(factor)
-
- output = tuple(sympy.Mul(*output[key][::-1])
- for key in ['lhs', 'operator', 'rhs'])
- return output
-
-
-def derivate(expression, operator):
- """ Calculate derivate of expression for given momentum operator:
-
- Parameters:
- -----------
- expression : sympy.Expr instance
- Sympy expression containing functions to to be derivated.
- operator : sympy.Symbol
- Sympy symbol representing momentum operator.
-
- Returns:
- --------
- output : sympy.Expr instance
- Derivated input expression.
- """
- if not isinstance(operator, sympy.Symbol):
- raise TypeError("Input operator '{}' is not type sympy.Symbol.")
-
- if operator.name not in ['k_x', 'k_y', 'k_z']:
- raise ValueError("Input operator '{}' unkown.".format(operator))
-
- if isinstance(expression, (int, float, sympy.Symbol)):
- return 0
- else:
- coordinate_name = operator.name.split('_')[1]
- ct = sympy.Symbol(coordinate_name, commutative=True)
- cf = sympy.Symbol(coordinate_name, commutative=False)
- h = sympy.Symbol('a_'+coordinate_name)
-
- expr1 = expression.subs({ct: ct + h, cf: cf + h})
- expr2 = expression.subs({ct: ct - h, cf: cf - h})
- output = (expr1 - expr2) / 2 / h
- return -sympy.I * sympy.expand(output)
-
-
-def _discretize_summand(summand, discrete_coordinates):
- """ Discretize one summand. """
- assert not isinstance(summand, sympy.Add), "Input should be one summand."
-
- def do_stuff(expr):
- """ Derivate expr recursively. """
- expr = sympy.expand(expr)
-
- if isinstance(expr, sympy.Add):
- return do_stuff(expr.args[-1]) + do_stuff(sympy.Add(*expr.args[:-1]))
-
- lhs, operator, rhs = split_factors(expr, discrete_coordinates)
- if rhs == 1 and operator != 1:
- return 0
- elif operator == 1:
- return lhs*rhs
- elif lhs == 1:
- return derivate(rhs, operator)
- else:
- return do_stuff(lhs*derivate(rhs, operator))
-
- return do_stuff(summand)
-
-
-def _discretize_expression(expression, discrete_coordinates):
- """ Discretize continous `expression` into discrete tb representation.
-
- Parameters:
- -----------
- expression : sympy.Expr instance
- The expression to be discretized.
-
- Returns:
- --------
- discrete_expression: dict
- dict in which key is offset of hopping ((0, 0, 0) for onsite)
- and value is corresponding symbolic hopping (onsite).
-
- Note:
- -----
- Recursive derivation implemented in _discretize_summand is applied
- on every summand. Shortening is applied before return on output.
- """
- if isinstance(expression, (int, float, sympy.Integer, sympy.Float)):
- n = len(discrete_coordinates)
- return {(tuple(0 for i in range(n))): expression}
-
- if not isinstance(expression, sympy.Expr):
- raise TypeError('Input expression should be a valid sympy expression.')
-
- coordinates_names = sorted(list(discrete_coordinates))
- coordinates = [sympy.Symbol(c, commutative=False) for c in coordinates_names]
- wf = sympy.Function(wavefunction_name)(*coordinates)
-
- if wf in expression.atoms(sympy.Function):
- raise ValueError("Input expression must not contain {}.".format(wf))
-
- expression = sympy.expand(expression*wf)
-
- # make sure we have list of summands
- summands = expression.args if expression.func == sympy.Add else [expression]
-
- # discretize every summand
- outputs = []
- for summand in summands:
- out = _discretize_summand(summand, discrete_coordinates)
- out = extract_hoppings(out, discrete_coordinates)
- outputs.append(out)
-
- # gather together
- discrete_expression = defaultdict(int)
- for summand in outputs:
- for k, v in summand.items():
- discrete_expression[k] += v
-
- return dict(discrete_expression)
-
-
-def discretize(hamiltonian, discrete_coordinates):
- """ Discretize continous `expression` into discrete tb representation.
-
- Parameters:
- -----------
- hamiltonian : sympy.Expr or sympy.Matrix instance
- The expression for the Hamiltonian.
-
- Returns:
- --------
- discrete_hamiltonian: dict
- dict in which key is offset of hopping ((0, 0, 0) for onsite)
- and value is corresponding symbolic hopping (onsite).
-
- Note:
- -----
- Recursive derivation implemented in _discretize_summand is applied
- on every summand. Shortening is applied before return on output.
- """
- if not isinstance(hamiltonian, sympy.matrices.MatrixBase):
- onsite_zeros = (0,)*len(discrete_coordinates)
- discrete_hamiltonian = {onsite_zeros: sympy.Integer(0)}
- hoppings = _discretize_expression(hamiltonian, discrete_coordinates)
- discrete_hamiltonian.update(hoppings)
- return discrete_hamiltonian
-
- shape = hamiltonian.shape
-
- discrete_hamiltonian = defaultdict(lambda: sympy.zeros(*shape))
- for i,j in itertools.product(range(shape[0]), range(shape[1])):
- expression = hamiltonian[i, j]
- hoppings = _discretize_expression(expression, discrete_coordinates)
-
- for offset, hop in hoppings.items():
- discrete_hamiltonian[offset][i,j] += hop
- return discrete_hamiltonian
-
-
-# ****** extracting hoppings ***********
-def read_hopping_from_wf(wf):
- """Read offset of a wave function in respect to (x,y,z).
-
- Parameters:
- ----------
- wf : sympy.function.AppliedUndef instance
- Function representing correct wave function used in discretizer.
- Should be created using global `wavefunction_name`.
-
- Returns:
- --------
- offset : tuple
- tuple of integers or floats that represent offset in respect to (x,y,z).
-
- Raises:
- -------
- ValueError
- If arguments of wf are repeated / do not stand for valid coordinates or
- lattice constants / order of dimensions is not lexical.
- TypeError:
- If wf is not of type sympy.function.AppliedUndef or its name does not
- corresponds to global 'wavefunction_name'.
- """
- if not isinstance(wf, sympy.function.AppliedUndef):
- raise TypeError('Input should be of type sympy.function.AppliedUndef.')
-
- if not wf.func.__name__ == wavefunction_name:
- msg = 'Input should be function that represents wavefunction in module.'
- raise TypeError(msg)
-
- # Check if input is consistent and in lexical order.
- # These are more checks for internal usage.
- coordinates_names = ['x', 'y', 'z']
- lattice_const_names = ['a_x', 'a_y', 'a_z']
- arg_coords = []
- for arg in wf.args:
- names = [s.name for s in arg.atoms(sympy.Symbol)]
- ind = -1
- for s in names:
- if not any(s in coordinates_names for s in names):
- raise ValueError("Wave function argument '{}' is incorrect.".format(s))
- if s not in coordinates_names and s not in lattice_const_names:
- raise ValueError("Wave function argument '{}' is incorrect.".format(s))
- if s in lattice_const_names:
- s = s.split('_')[1]
- tmp = coordinates_names.index(s)
- if tmp in arg_coords:
- msg = "Wave function '{}' arguments are inconsistent."
- raise ValueError(msg.format(wf))
- if ind != -1:
- if tmp != ind:
- msg = "Wave function '{}' arguments are inconsistent."
- raise ValueError(msg.format(wf))
- else:
- ind = tmp
- arg_coords.append(ind)
- if arg_coords != sorted(arg_coords):
- msg = "Coordinates of wave function '{}' are not in lexical order."
- raise ValueError(msg.format(wf))
-
- # function real body
- offset = []
- for argument in wf.args:
- temp = sympy.expand(argument)
- if temp in sympy.symbols('x y z', commutative = False):
- offset.append(0)
- elif temp.func == sympy.Add:
- for arg_summands in temp.args:
- if arg_summands.func == sympy.Mul:
- if len(arg_summands.args) > 2:
- print('More than two factors in an argument of wf')
- if not arg_summands.args[0] in sympy.symbols('a_x a_y a_z'):
- offset.append(arg_summands.args[0])
- else:
- offset.append(arg_summands.args[1])
- elif arg_summands in sympy.symbols('a_x a_y a_z'):
- offset.append(1)
- else:
- print('Argument of \wf is neither a sum nor a single space variable.')
- return tuple(offset)
-
-
-def extract_hoppings(expression, discrete_coordinates):
- """Extract hopping and perform shortening operation. """
-
- # make sure we have list of summands
- expression = sympy.expand(expression)
- summands = expression.args if expression.func == sympy.Add else [expression]
-
- hoppings = defaultdict(int)
- for summand in summands:
- if summand.func.__name__ == wavefunction_name:
- hoppings[read_hopping_from_wf(summand)] += 1
- else:
- for i in range(len(summand.args)):
- if summand.args[i].func.__name__ == wavefunction_name:
- index = i
- if index < len(summand.args) - 1:
- print('Psi is not in the very end of the term. Output will be wrong!')
- hoppings[read_hopping_from_wf(summand.args[-1])] += sympy.Mul(*summand.args[:-1])
-
- # START: shortenig
- discrete_coordinates = sorted(list(discrete_coordinates))
- tmps = ['a_{}'.format(s) for s in discrete_coordinates]
- lattice_constants = sympy.symbols(tmps)
- a = sympy.Symbol('a')
-
- # make a list of all hopping kinds we have to consider during the shortening
- hops_kinds = np.array(list(hoppings))
- # find the longest hopping range in each direction
- longest_ranges = [np.max(hops_kinds[:,i]) for i in range(len(hops_kinds[0,:]))]
- # define an array in which we are going to store by which factor we
- # can shorten the hoppings in each direction
- shortening_factors = np.ones_like(longest_ranges)
- # Loop over the direction and each potential shortening factor.
- # Inside the loop test whether the hopping distances are actually
- # multiples of the potential shortening factor.
- for dim in np.arange(len(longest_ranges)):
- for factor in np.arange(longest_ranges[dim])+1:
- modulos = np.mod(hops_kinds[:, dim], factor)
- if np.sum(modulos) < 0.1:
- shortening_factors[dim] = factor
- # Apply the shortening factors on the hopping.
- short_hopping = {}
- for hopping_kind in hoppings.keys():
- short_hopping_kind = tuple(np.array(hopping_kind) / shortening_factors)
-
- for i in short_hopping_kind:
- if isinstance(i, float):
- assert i.is_integer()
- short_hopping_kind = tuple(int(i) for i in short_hopping_kind)
-
- short_hopping[short_hopping_kind] = hoppings[hopping_kind]
- for lat_const, factor in zip(lattice_constants, shortening_factors):
- factor = int(factor)
- subs = {lat_const: lat_const/factor}
- short_hopping[short_hopping_kind] = short_hopping[short_hopping_kind].subs(subs)
-
- # We don't need separate a_x, a_y and a_z anymore.
- for key, val in short_hopping.items():
- short_hopping[key] = val.subs({i: a for i in sympy.symbols('a_x a_y a_z')})
-
- return short_hopping
diff --git a/discretizer/discretizer.py b/discretizer/discretizer.py
deleted file mode 100644
index 70b2a9dac8368a0f00413e018051b2c4527fa594..0000000000000000000000000000000000000000
--- a/discretizer/discretizer.py
+++ /dev/null
@@ -1,160 +0,0 @@
-from __future__ import print_function, division
-
-import warnings
-import numpy as np
-import sympy
-
-from .algorithms import read_coordinates
-from .algorithms import discretize
-
-from .postprocessing import offset_to_direction
-from .postprocessing import make_kwant_functions
-from .postprocessing import offset_to_direction
-
-from .interpolation import interpolate_tb_hamiltonian
-
-try:
- # normal situation
- from kwant import Builder
- from kwant import TranslationalSymmetry
- from kwant import HoppingKind
- from kwant.lattice import Monatomic
-except ImportError:
- # probably run on gitlab-ci
- pass
-
-class Discretizer(object):
- """Discretize continous Hamiltonian into its tight binding representation.
-
- This class provides easy and nice interface for passing models to Kwant.
-
- Parameters:
- -----------
- hamiltonian : sympy.Expr or sympy.Matrix instance
- Symbolic representation of a continous Hamiltonian. Momentum operators
- should be taken from ``discretizer.momentum_operators``.
- discrete_coordinates : set of strings
- Set of coordinates for which momentum operators will be treated as
- differential operators. For example ``discrete_coordinates={'x', 'y'}``.
- If left as a None they will be obtained from the input hamiltonian by
- reading present coordinates and momentum operators.
- interpolate : bool
- If True all space dependent parameters in onsite and hopping will be
- interpolated to depenend only on the values at site positions.
- Default is False.
- both_hoppings_directions : bool
- If True all hoppings will be returned. For example, if set to True, both
- hoppings into (1, 0) and (-1, 0) will be returned. Default is False.
- verbose : bool
- If True additional information will be printed. Default is False.
-
- Attributes:
- -----------
- symbolic_hamiltonian : dictionary
- Dictionary containing symbolic result of discretization. Key is the
- direction of the hopping (zeros for onsite)
- lattice : kwant.lattice.Monatomic instance
- Lattice to create kwant system. Lattice constant is set to
- lattice_constant value.
- onsite : function
- The value of the onsite Hamiltonian.
- hoppings : dict
- A dictionary with keys being tuples of the lattice hopping, and values
- the corresponding value functions.
- discrete_coordinates : set of strings
- As in input.
- input_hamiltonian : sympy.Expr or sympy.Matrix instance
- The input hamiltonian after preprocessing (substitution of functions).
- """
- def __init__(self, hamiltonian, discrete_coordinates=None,
- lattice_constant=1, interpolate=False,
- both_hoppings_directions=False, verbose=False):
-
- self.input_hamiltonian = hamiltonian
-
- if discrete_coordinates is None:
- self.discrete_coordinates = read_coordinates(hamiltonian)
- else:
- self.discrete_coordinates = discrete_coordinates
-
- if verbose:
- print('Discrete coordinates set to: ',
- sorted(self.discrete_coordinates), end='\n\n')
-
- # discretization
- if self.discrete_coordinates:
- tb_ham = discretize(hamiltonian, self.discrete_coordinates)
- tb_ham = offset_to_direction(tb_ham, self.discrete_coordinates)
- else:
- tb_ham = {(0,0,0): hamiltonian}
- self.discrete_coordinates = {'x', 'y', 'z'}
-
- if interpolate:
- tb_ham = interpolate_tb_hamiltonian(tb_ham)
-
- if not both_hoppings_directions:
- keys = list(tb_ham)
- tb_ham = {k: v for k, v in tb_ham.items()
- if k in sorted(keys)[len(keys)//2:]}
-
- self.symbolic_hamiltonian = tb_ham.copy()
-
- for key, val in tb_ham.items():
- tb_ham[key] = val.subs(sympy.Symbol('a'), lattice_constant)
-
- # making kwant lattice
- dim = len(self.discrete_coordinates)
-
- self.lattice = Monatomic(lattice_constant*np.eye(dim).reshape(dim, dim))
- self.lattice_constant = lattice_constant
-
- # making kwant functions
- tb = make_kwant_functions(tb_ham, self.discrete_coordinates, verbose)
- self.onsite = tb.pop((0,)*len(self.discrete_coordinates))
- self.hoppings = {HoppingKind(d, self.lattice): val
- for d, val in tb.items()}
-
- def build(self, shape, start, symmetry=None, periods=None):
- """Build Kwant's system.
-
- Convienient functions that simplifies building of a Kwant's system.
-
- Parameters:
- -----------
- shape : function
- A function of real space coordinates that returns a truth value:
- true for coordinates inside the shape, and false otherwise.
- start : 1d array-like
- The real-space origin for the flood-fill algorithm.
- symmetry : 1d array-like
- Deprecated. Please use ```periods=[symmetry]`` instead.
- periods : list of tuples
- If periods are provided a translational invariant system will be
- built. Periods corresponds basically to a translational symmetry
- defined in real space. This vector will be scalled by a lattice
- constant before passing it to ``kwant.TranslationalSymmetry``.
- Examples: ``periods=[(1,0,0)]`` or ``periods=[(1,0), (0,1)]``.
- In second case one will need https://gitlab.kwant-project.org/cwg/wraparound
- in order to finalize system.
-
- Returns:
- --------
- system : kwant.Builder instance
- """
- if symmetry is not None:
- warnings.warn("\nSymmetry argument is deprecated. " +
- "Please use ```periods=[symmetry]`` instead.",
- DeprecationWarning)
- periods = [symmetry]
-
- if periods is None:
- sys = Builder()
- else:
- vecs = [self.lattice.vec(p) for p in periods]
- sys = Builder(TranslationalSymmetry(*vecs))
-
- sys[self.lattice.shape(shape, start)] = self.onsite
- for hop, val in self.hoppings.items():
- sys[hop] = val
-
- return sys
diff --git a/discretizer/interpolation.py b/discretizer/interpolation.py
deleted file mode 100644
index da83f0ace11cfada05db549c161f4ec05dd0e759..0000000000000000000000000000000000000000
--- a/discretizer/interpolation.py
+++ /dev/null
@@ -1,109 +0,0 @@
-from __future__ import print_function, division
-
-import itertools
-import numpy as np
-import sympy
-
-
-def _follow_path(expr, path):
- res = expr
- for i in np.arange(len(path)):
- res = res.args[path[i]]
- return res
-
-def _interchange(expr, sub, path):
- res = sub
- for i in np.arange(len(path)):
- temp = _follow_path(expr, path[:-(i+1)])
- args = list(temp.args)
- args[path[len(path)-i-1]] = res
- res = temp.func(*tuple(args))
- return res
-
-def _interpolate_Function(expr):
- path = 'None'
- factor = 'None'
- change = False
- summand_0 = 'None'
- summand_1 = 'None'
- res = expr
- for i in np.arange(len(expr.args)):
- argument = sympy.expand(expr.args[i])
- if argument.func == sympy.Add:
- for j in np.arange(len(argument.args)):
- summand = argument.args[j]
- if summand.func == sympy.Mul:
- for k in np.arange(len(summand.args)):
- temp = 0
- if summand.args[k] == sympy.Symbol('a'):
- temp = sympy.Mul(sympy.Mul(*summand.args[:k]),
- sympy.Mul(*summand.args[k+1:]))
- #print(temp)
- if not temp == int(temp):
- #print('found one')
- factor = (temp)
- path = np.array([i, j, k])
- if not factor == 'None':
- change = True
-
- sign = np.sign(factor)
- offsets = np.array([int(factor), sign * (int(sign * factor) + 1)])
- weights = 1/np.abs(offsets - factor)
- weights = weights/np.sum(weights)
-
- res = ( weights[0] * _interchange(expr, offsets[0] * sympy.Symbol('a'), path[:-1])
- + weights[1] * _interchange(expr, offsets[1] * sympy.Symbol('a'), path[:-1]))
-
- return sympy.expand(res), change
-
-
-def _interpolate_expression(expr):
- change = False
- expr = sympy.expand(expr)
- res = expr
-
- if isinstance(expr, sympy.Function):# and not change:
- path = np.array([])
- temp, change = _interpolate_Function(_follow_path(expr, path))
- res = _interchange(expr, temp, path)
-
- for i in np.arange(len(expr.args)):
- path = np.array([i])
- if isinstance(_follow_path(expr, path), sympy.Function) and not change:
- temp, change = _interpolate_Function(_follow_path(expr, path))
- res = _interchange(expr, temp, path)
-
- for j in np.arange(len(expr.args[i].args)):
- path = np.array([i, j])
- if isinstance(_follow_path(expr, path), sympy.Function) and not change:
- temp, change = _interpolate_Function(_follow_path(expr, path))
- res = _interchange(expr, temp, path)
-
- if change:
- res = _interpolate_expression(res)
-
- return sympy.expand(res)
-
-
-def _interpolate(expression):
- if not isinstance(expression, sympy.Matrix):
- return _interpolate_expression(expression)
-
- shape = expression.shape
- interpolated = sympy.zeros(*shape)
- for i,j in itertools.product(range(shape[0]), repeat=2):
- interpolated[i,j] = _interpolate_expression(expression[i, j])
-
- return interpolated
-
-
-def interpolate_tb_hamiltonian(tb_hamiltonian):
- """Interpolate tight binding hamiltonian.
-
- This function perform linear interpolation to provide onsite and hoppings
- depending only on parameters values at sites positions.
- """
- interpolated = {}
- for key, val in tb_hamiltonian.items():
- interpolated[key] = _interpolate(val)
- return interpolated
diff --git a/discretizer/postprocessing.py b/discretizer/postprocessing.py
deleted file mode 100644
index cd2fe61a0e5e9209852375345d9fd261350cf815..0000000000000000000000000000000000000000
--- a/discretizer/postprocessing.py
+++ /dev/null
@@ -1,214 +0,0 @@
-from __future__ import print_function, division
-
-import sympy
-from sympy.utilities.lambdify import lambdastr
-from sympy.printing.lambdarepr import LambdaPrinter
-from sympy.core.function import AppliedUndef
-
-
-class NumericPrinter(LambdaPrinter):
- def _print_ImaginaryUnit(self, expr):
- return "1.j"
-
-
-def offset_to_direction(discrete_hamiltonian, discrete_coordinates):
- """Translate hopping keys from offsets to directions.
-
- Parameters:
- -----------
- discrete_hamiltonian: dict
- Discretized hamiltonian, key should be an offset of a hopping and value
- corresponds to symbolic hopping.
- discrete_coordinates: set
- Set of discrete coordinates.
-
- Returns:
- --------
- discrete_hamiltonian: dict
- Discretized hamiltonian, key is a direction of a hopping and value
- corresponds to symbolic hopping.
-
- Note:
- -----
- Coordinates (x,y,z) in output stands for a position of a source of the
- hopping.
- """
- coordinates = sorted(list(discrete_coordinates))
- coordinates = [sympy.Symbol(s, commutative=False) for s in coordinates]
- a = sympy.Symbol('a')
-
- onsite_zeros = (0,)*len(discrete_coordinates)
- output = {onsite_zeros: discrete_hamiltonian.pop(onsite_zeros)}
- for offset, hopping in discrete_hamiltonian.items():
- direction = tuple(-c for c in offset)
- subs = {c: c + d*a for c, d in zip(coordinates, direction)}
- output[direction] = hopping.subs(subs)
-
- return output
-
-# ************ Making kwant functions ***********
-def make_return_string(expr):
- """Process a sympy expression into an evaluatable Python return statement.
-
- Parameters:
- -----------
- expr : sympy.Expr instance
-
- Returns:
- --------
- output : string
- A return string that can be used to assemble a Kwant value function.
- func_symbols : set of sympy.Symbol instances
- All space dependent functions that appear in the expression.
- const_symbols : set of sympy.Symbol instances
- All constants that appear in the expression.
- """
- func_symbols = {sympy.Symbol(i.func.__name__) for i in
- expr.atoms(AppliedUndef)}
-
- free_symbols = {i for i in expr.free_symbols if i.name not in ['x', 'y', 'z']}
- const_symbols = free_symbols - func_symbols
-
- expr = expr.subs(sympy.I, sympy.Symbol('1.j')) # quick hack
- output = lambdastr((), expr, printer=NumericPrinter)[len('lambda : '):]
- output = output.replace('MutableDenseMatrix', 'np.array')
- output = output.replace('ImmutableMatrix', 'np.array')
-
- return 'return {}'.format(output), func_symbols, const_symbols
-
-
-def assign_symbols(func_symbols, const_symbols, discrete_coordinates,
- onsite=True):
- """Generate a series of assingments defining a set of symbols.
-
- Parameters:
- -----------
- func_symbols : set of sympy.Symbol instances
- All space dependent functions that appear in the expression.
- const_symbols : set of sympy.Symbol instances
- All constants that appear in the expression.
-
- Returns:
- --------
- assignments : list of strings
- List of lines used for including in a function.
-
- Notes:
- where A, B, C are all the free symbols plus the symbols that appear on the
- ------
- The resulting lines begin with a coordinates assignment of form
- `x,y,z = site.pos` when onsite=True, or
- `x,y,z = site2.pos` when onsite=False
-
- followed by two lines of form
- `A, B, C = p.A, p.B, p.C`
- `f, g, h = p.f, p.g, p.h`
- where A, B, C are symbols representing constants and f, g, h are symbols
- representing functions. Separation of constant and func symbols is probably
- not necessary but I leave it for now, just in case.
- """
- lines = []
- func_names = [i.name for i in func_symbols]
- const_names = [i.name for i in const_symbols]
-
- if func_names:
- lines.insert(0, ', '.join(func_names) + ' = p.' +
- ', p.'.join(func_names))
-
- if const_names:
- lines.insert(0, ', '.join(const_names) + ' = p.' +
- ', p.'.join(const_names))
-
- if onsite:
- site = 'site'
- else:
- site = 'site2'
-
- names = sorted(list(discrete_coordinates))
- lines.insert(0, '({}, ) = {}.pos'.format(', '.join(names), site))
-
- return lines
-
-
-def value_function(content, name='_anonymous_func', onsite=True, verbose=False):
- """Generate a Kwant value function from a list of lines containing its body.
-
- Parameters:
- -----------
- content : list of lines
- Lines forming the body of the function.
- name : string
- Function name (not important).
- onsite : bool
- If True, the function call signature will be `f(site, p)`, otherwise
- `f(site1, site2, p)`.
- verbose : bool
- Whether the function bodies should be printed.
-
- Returns:
- --------
- f : function
- The function defined in a separated namespace.
- """
- if not content[-1].startswith('return'):
- raise ValueError('The function does not end with a return statement')
-
- separator = '\n' + 4 * ' '
- site_string = 'site' if onsite else 'site1, site2'
- header = 'def {0}({1}, p):'.format(name, site_string)
- func_code = separator.join([header] + list(content))
-
- namespace = {}
- if verbose:
- print(func_code)
- exec("from __future__ import division", namespace)
- exec("import numpy as np", namespace)
- exec("from numpy import *", namespace)
- exec(func_code, namespace)
- return namespace[name]
-
-
-def make_kwant_functions(discrete_hamiltonian, discrete_coordinates,
- verbose=False):
- """Transform discrete hamiltonian into valid kwant functions.
-
- Parameters:
- -----------
- discrete_hamiltonian: dict
- dict in which key is offset of hopping ((0, 0, 0) for onsite)
- and value is corresponding symbolic hopping (onsite).
-
- verbose : bool
- Whether the function bodies should be printed.
-
- discrete_coordinates : tuple/list
- List of discrete coordinates. Must corresponds to offsets in
- discrete_hamiltonian keys.
-
- Note:
- -----
-
- """
- dim = len(discrete_coordinates)
- if not all(len(i)==dim for i in list(discrete_hamiltonian.keys())):
- raise ValueError("Dimension of offsets and discrete_coordinates" +
- "do not match.")
-
- functions = {}
- for offset, hopping in discrete_hamiltonian.items():
- onsite = True if all(i == 0 for i in offset) else False
- return_string, func_symbols, const_symbols = make_return_string(hopping)
- lines = assign_symbols(func_symbols, const_symbols, onsite=onsite,
- discrete_coordinates=discrete_coordinates)
- lines.append(return_string)
-
- if verbose:
- print("Function generated for {}:".format(offset))
- f = value_function(lines, verbose=verbose, onsite=onsite)
- print()
- else:
- f = value_function(lines, verbose=verbose, onsite=onsite)
-
- functions[offset] = f
-
- return functions
diff --git a/discretizer/tests/test_algorithms.py b/discretizer/tests/test_algorithms.py
deleted file mode 100644
index cd253183cee947de67e1ef724e7ec6db7f1c7e7b..0000000000000000000000000000000000000000
--- a/discretizer/tests/test_algorithms.py
+++ /dev/null
@@ -1,120 +0,0 @@
-from __future__ import print_function, division
-
-import sympy
-import discretizer
-from discretizer.algorithms import read_coordinates
-from discretizer.algorithms import split_factors
-from discretizer.algorithms import wavefunction_name
-from discretizer.algorithms import derivate
-from discretizer.algorithms import _discretize_summand
-
-from nose.tools import raises
-from nose.tools import assert_raises
-import numpy as np
-
-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')
-
-wf = sympy.Function(wavefunction_name)
-Psi = sympy.Function(wavefunction_name)(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_read_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', 'y'},
- kx**2 + kz * B(y) : {'x', 'y', 'z'},
- }
- for inp, out in test.items():
- got = read_coordinates(inp)
- assert got == out,\
- "Should be: split_factors({})=={}. Not {}".format(inp, out, got)
-
-def test_split_factors_1():
- test = {
- kz * Psi : (1, kz, Psi),
- A * kx**2 * Psi : (A * kx, kx, Psi),
- A * kx**2 * ky * Psi : (A * kx**2, ky, Psi),
- ky * A * kx * B * Psi : (ky * A, kx, B * Psi),
- kx : (1, kx, 1),
- kx**2 : (kx, kx, 1),
- A : (1, 1, A),
- A**2 : (1, 1, A**2),
- kx*A**2 : (1, kx, A**2),
- kx**2*A**2 : (kx, kx, A**2),
- A(x, y, z) : (1, 1, A(x, y, z)),
- Psi : (1, 1, Psi),
- np.int(5) : (1, 1, np.int(5)),
- np.float(5) : (1, 1, np.float(5)),
- sympy.Integer(5) : (1, 1, sympy.Integer(5)),
- sympy.Float(5) : (1, 1, sympy.Float(5)),
- 1 : (1, 1, 1),
- 1.0 : (1, 1, 1.0),
- 5 : (1, 1, 5),
- 5.0 : (1, 1, 5.0),
- }
-
- for inp, out in test.items():
- got = split_factors(inp, discrete_coordinates={'x', 'y', 'z'})
- assert got == out,\
- "Should be: split_factors({})=={}. Not {}".format(inp, out, got)
-
-
-@raises(AssertionError)
-def test_split_factors_2():
- split_factors(A+B, discrete_coordinates={'x', 'y', 'z'})
-
-
-def test_derivate_1():
- test = {
- (A(x), kx): '-I*(-A(-a_x + x)/(2*a_x) + A(a_x + x)/(2*a_x))',
- (A(x), ky): '0',
- (A(x)*B, kx): '-I*(-A(-a_x + x)*B/(2*a_x) + A(a_x + x)*B/(2*a_x))',
- (A(x) + B(x), kx): '-I*(-A(-a_x + x)/(2*a_x) + A(a_x + x)/(2*a_x) - B(-a_x + x)/(2*a_x) + B(a_x + x)/(2*a_x))',
- (A, kx): '0',
- (5, kx): '0',
- (A(x) * B(x), kx): '-I*(-A(-a_x + x)*B(-a_x + x)/(2*a_x) + A(a_x + x)*B(a_x + x)/(2*a_x))',
- (Psi, ky): '-I*(-Psi(x, -a_y + y, z)/(2*a_y) + Psi(x, a_y + y, z)/(2*a_y))',
- }
-
- for inp, out in test.items():
- got = (derivate(*inp))
- out = sympy.sympify(out, locals=ns)
- assert sympy.simplify(sympy.expand(got - out)) == 0,\
- "Should be: derivate({0[0]}, {0[1]})=={1}. Not {2}".format(inp, out, got)
-
-
-@raises(TypeError)
-def test_derivate_2():
- derivate(A(x), kx**2)
-
-
-@raises(ValueError)
-def test_derivate_3():
- derivate(A(x), sympy.Symbol('A'))
-
-
-def test_discretize_summand_1():
- test = {
- kx * A(x): '-I*(-A(-a_x + x)/(2*a_x) + A(a_x + x)/(2*a_x))',
- kx * Psi: '-I*(-Psi(-a_x + x, y, z)/(2*a_x) + Psi(a_x + x, y, z)/(2*a_x))',
- kx**2 * Psi: 'Psi(x, y, z)/(2*a_x**2) - Psi(-2*a_x + x, y, z)/(4*a_x**2) - Psi(2*a_x + x, y, z)/(4*a_x**2)',
- kx * A(x) * kx * Psi: 'A(-a_x + x)*Psi(x, y, z)/(4*a_x**2) - A(-a_x + x)*Psi(-2*a_x + x, y, z)/(4*a_x**2) + A(a_x + x)*Psi(x, y, z)/(4*a_x**2) - A(a_x + x)*Psi(2*a_x + x, y, z)/(4*a_x**2)',
- }
-
- for inp, out in test.items():
- got = _discretize_summand(inp, discrete_coordinates={'x', 'y', 'z'})
- out = sympy.sympify(out, locals=ns)
- assert sympy.simplify(sympy.expand(got - out)) == 0,\
- "Should be: _discretize_summand({})=={}. Not {}".format(inp, out, got)
-
-@raises(AssertionError)
-def test_discretize_summand_2():
- _discretize_summand(kx*A(x)+ B(x), discrete_coordinates={'x', 'y', 'z'})
diff --git a/discretizer/tests/test_hoppings.py b/discretizer/tests/test_hoppings.py
deleted file mode 100644
index 21075c70ad724294a5c2672a2ff8f290cd5dde30..0000000000000000000000000000000000000000
--- a/discretizer/tests/test_hoppings.py
+++ /dev/null
@@ -1,129 +0,0 @@
-from __future__ import print_function, division
-
-import sympy
-import discretizer
-from discretizer.algorithms import wavefunction_name
-from discretizer.algorithms import read_hopping_from_wf
-from discretizer.algorithms import extract_hoppings
-
-from nose.tools import raises
-from nose.tools import assert_raises
-import numpy as np
-
-
-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')
-
-wf = sympy.Function(wavefunction_name)
-Psi = sympy.Function(wavefunction_name)(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_read_hoppings_from_wf_1():
- offsets = [(0,0,0), (1,0,0), (0,1,0), (0,0,1), (1,1,1), (1,2,3)]
- test = {}
-
- for offset in offsets:
- nx, ny, nz = offset
- key = Psi.subs({x: x + nx*ax, y: y + ny * ay, z: z + nz * az})
- test[key] = offset
-
- for inp, out in test.items():
- got = read_hopping_from_wf(inp)
- assert got == out,\
- "Should be: read_hopping_from_wf({}) == {}. Not {}".format(inp, out, got)
-
-
-def test_read_hoppings_from_wf_2():
- test = {
- wf(x, y, z): (0,0,0),
- wf(x, y): (0, 0),
- wf(x, z): (0, 0),
- wf(x+ax, y-2*ay): (1, -2),
- wf(x, z+3*az): (0, 3),
- wf(y, z): (0, 0),
- wf(y, z+az): (0, 1),
- }
-
- for inp, out in test.items():
- got = read_hopping_from_wf(inp)
- assert got == out,\
- "Should be: read_hopping_from_wf({}) == {}. Not {}".format(inp, out, got)
-
-
-def test_test_read_hoppings_from_wf_ValueError():
- tests = {
- wf(x+ay, ay),
- wf(y, x),
- wf(x, x),
- wf(x, ax),
- wf(y, A),
- }
- for inp in tests:
- assert_raises(ValueError, read_hopping_from_wf, inp)
-
-
-def test_test_read_hoppings_from_wf_TypeError():
- tests = {
- wf(x,y,z) + A,
- A(x,y),
- 5*Psi,
- Psi+2,
- A,
- }
- for inp in tests:
- assert_raises(TypeError, read_hopping_from_wf, inp)
-
-
-
-def test_extract_hoppings():
- discrete_coordinates = {'x', 'y', 'z'}
- tests = [
- {
- 'test_inp': '-I*(-Psi(-a_x + x, y)/(2*a_x) + Psi(a_x + x, y)/(2*a_x))',
- 'test_out': {
- (1, 0): '-I/(2*a)',
- (-1, 0): 'I/(2*a)',
- },
- },
- {
- 'test_inp': 'Psi(x, y)/(2*a_x**2) - Psi(-2*a_x + x, y)/(4*a_x**2) - Psi(2*a_x + x, y)/(4*a_x**2)',
- 'test_out': {
- (1, 0): '-1/a**2',
- (0, 0): '2/a**2',
- (-1, 0): '-1/a**2',
- },
- },
- {
- 'test_inp': 'A(-a_x + x)*Psi(x, y)/(4*a_x**2) - A(-a_x + x)*Psi(-2*a_x + x, y)/(4*a_x**2) + A(a_x + x)*Psi(x, y)/(4*a_x**2) - A(a_x + x)*Psi(2*a_x + x, y)/(4*a_x**2)',
- 'test_out': {
- (1, 0): '-A(a/2 + x)/a**2',
- (0, 0): 'A(-a/2 + x)/a**2 + A(a/2 + x)/a**2',
- (-1, 0): '-A(-a/2 + x)/a**2',
- },
- },
- {
- 'test_inp': 'I*A(x, y)*Psi(x, -a_y + y, z)/(4*a_x**2*a_y) - I*A(x, y)*Psi(x, a_y + y, z)/(4*a_x**2*a_y) - I*A(-2*a_x + x, y)*Psi(-2*a_x + x, -a_y + y, z)/(8*a_x**2*a_y) + I*A(-2*a_x + x, y)*Psi(-2*a_x + x, a_y + y, z)/(8*a_x**2*a_y) - I*A(2*a_x + x, y)*Psi(2*a_x + x, -a_y + y, z)/(8*a_x**2*a_y) + I*A(2*a_x + x, y)*Psi(2*a_x + x, a_y + y, z)/(8*a_x**2*a_y)',
- 'test_out': {
- (-1, 1, 0): 'I*A(-a + x, y)/(2*a**3)',
- (1, 1, 0): 'I*A(a + x, y)/(2*a**3)',
- (0, -1, 0): 'I*A(x, y)/a**3',
- (0, 1, 0): '-I*A(x, y)/a**3',
- (1, -1, 0): '-I*A(a + x, y)/(2*a**3)',
- (-1, -1, 0): '-I*A(-a + x, y)/(2*a**3)',
- },
- },
- ]
-
- for test in tests:
- test_inp = test['test_inp']
- test_out = test['test_out']
- inp = sympy.sympify(test_inp, locals=ns)
- result = extract_hoppings(inp, discrete_coordinates)
- for key, got in result.items():
- out = sympy.sympify(test_out[key], locals=ns)
- assert sympy.simplify(sympy.expand(got - out)) == 0, \
- "Should be: extract_hoppings({})=={}. Not {}".format(inp, test_out, result)
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 == ''