merge discretizer into Kwant

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

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)

kwant/continuum/__init__.py

+__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

kwant/continuum/_common.py

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

kwant/continuum/tests/test_common.py

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

kwant/continuum/tests/test_discretizer.py

+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 == ''