improve usage of 'substitutions' in 'continuum'

So far substitutions were accepted in form ''string key: sympy expression'',
which was extended to accept form of ''string key: string expression'.
Also 'substitutions' were only used when input 'hamiltonian' was not a sympy
object already, which got alternated to apply them always. Closes #97.

kwant/continuum/_common.py

@@ -48,13 +48,23 @@ def lambdify(hamiltonian, *, substitutions=None):
         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]]})``.
+        A namespace of substitutions to be performed on the input
+        ``hamiltonian``. It can be used to simplify input of matrices or
+        alternate input before proceeding further. Please see examples below.
+
+    Example:
+    --------
+        >>> f = kwant.continuum.lambdify('a + b', substitutions={'b': 'b + c'})
+        >>> f(1, 3, 5)
+        9
+
+        >>> subs = {'s_z': [[1, 0], [0, -1]]}
+        >>> f = kwant.continuum.lambdify('k_z**2 * s_z', substitutions=subs)
+        >>> f(.25)
+        array([[ 0.0625,  0.    ],
+               [ 0.    , -0.0625]])
     """
-    expr = hamiltonian
-    if not isinstance(expr, (sympy.Expr, sympy.matrices.MatrixBase)):
-        expr = sympify(expr, substitutions)
+    expr = sympify(hamiltonian, substitutions)
 
     args = [s.name for s in expr.atoms(sympy.Symbol)]
     args += [str(f.func) for f in expr.atoms(AppliedUndef, sympy.Function)]
@@ -68,7 +78,7 @@ def lambdify(hamiltonian, *, substitutions=None):
     return f
 
 
-def sympify(string, substitutions=None, **kwargs):
+def sympify(e, substitutions=None):
     """Return sympified object with respect to kwant-specific rules.
 
     This is a modification of ``sympy.sympify`` to apply kwant-specific rules,
@@ -77,37 +87,63 @@ def sympify(string, substitutions=None, **kwargs):
 
     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.
+    e : arbitrary expression
+        An expression that will be converted to a sympy object.
+        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:
+        A namespace of substitutions to be performed on the input ``e``.
+        It works in a similar way to ``locals`` argument in ``sympy.sympify``
+        but extends its functionality to i.e. 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``.
+        Keys should be strings, unless ``e`` is already a sympy object.
+        Values can be strings or ``sympy`` objects.
 
     Example
     -------
         >>> from kwant.continuum import sympify
-        >>> hamiltonian = sympify('k_x * A(x) * k_x + V(x)')
+        >>> sympify('k_x * A(x) * k_x + V(x)')
+        k_x*A(x)*k_x + V(x)     # as valid sympy object
+
+        or
+
+        >>> from kwant.continuum import sympify
+        >>> sympify('k_x**2 + V', substitutions={'V': 'V_0 + V(x)'})
+        k_x**2 + V(x) + V_0
     """
     stored_value = None
-
+    sympified_types = (sympy.Expr, sympy.matrices.MatrixBase)
     if substitutions is None:
         substitutions = {}
 
-    substitutions.update(_clash)
-
+    # if ``e`` is already a ``sympy`` object we make use of ``substitutions``
+    # and terminate a code path.
+    if isinstance(e, sympified_types):
+        subs = {sympify(k): sympify(v) for k, v in substitutions.items()}
+        return e.subs(subs)
+
+    # if ``e`` is not a sympified type then we proceed with sympifying process,
+    # we expect all keys in ``substitutions`` to be strings at this moment.
+    if not all(isinstance(k, str) for k in substitutions):
+        raise ValueError("If 'e' is not already a sympy object ",
+                         "then keys of 'substitutions' must be strings.")
+
+    # sympify values of substitutions before updating it with _clash
+    substitutions = {k: (sympify(v) if not isinstance(v, sympified_types)
+                         else v)
+                     for k, v in substitutions.items()}
+
+    substitutions.update({s: v for s, v in _clash.items()
+                          if s not in substitutions})
     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)
+        hamiltonian = sympy.sympify(e, locals=substitutions)
+        # if input is for example ``[[k_x * A(x) * k_x]]`` after the first
+        # sympify we are getting list of sympy objects, so we call sympify
+        # second time to obtain ``sympy`` matrices.
+        hamiltonian = sympy.sympify(hamiltonian)
     finally:
         if stored_value is not None:
             converter[list] = stored_value

kwant/continuum/discretizer.py

@@ -49,9 +49,12 @@ def discretize(hamiltonian, discrete_coordinates=None, lattice_constant=1,
     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]]})``.
+        A namespace of substitutions to be performed on the input
+        ``hamiltonian``. Values can be either strings or ``sympy`` objects.
+        It can be used to simplify input of matrices or alternate input before
+        proceeding further. For example:
+        ``substitutions={'k': 'k_x + I * k_y'}`` or
+        ``substitutions={'s_z': [[1, 0], [0, -1]]}``.
     verbose : bool, default: False
         If ``True`` additional information will be printed.
 
@@ -85,9 +88,12 @@ def discretize_symbolic(hamiltonian, discrete_coordinates=None, substitutions=No
         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]]})``.
+        A namespace of substitutions to be performed on the input
+        ``hamiltonian``. Values can be either strings or ``sympy`` objects.
+        It can be used to simplify input of matrices or alternate input before
+        proceeding further. For example:
+        ``substitutions={'k': 'k_x + I * k_y'}`` or
+        ``substitutions={'s_z': [[1, 0], [0, -1]]}``.
     verbose : bool, default: False
         If ``True`` additional information will be printed.
 
@@ -101,8 +107,7 @@ def discretize_symbolic(hamiltonian, discrete_coordinates=None, substitutions=No
         discrete_coordinates : sequence of strings
             The coordinates that have been discretized.
     """
-    if not isinstance(hamiltonian, (sympy.Expr, sympy.matrices.MatrixBase)):
-        hamiltonian = sympify(hamiltonian, substitutions)
+    hamiltonian = sympify(hamiltonian, substitutions)
 
     atoms_names = [s.name for s in hamiltonian.atoms(sympy.Symbol)]
     if any( s == 'a' for s in atoms_names):
@@ -176,9 +181,12 @@ def build_discretized(tb_hamiltonian, discrete_coordinates,
     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]]})``.
+        A namespace of substitutions to be performed on the values of input
+        ``tb_hamiltonian``. Values can be either strings or ``sympy`` objects.
+        It can be used to simplify input of matrices or alternate input before
+        proceeding further. For example:
+        ``substitutions={'k': 'k_x + I * k_y'}`` or
+        ``substitutions={'s_z': [[1, 0], [0, -1]]}``.
     verbose : bool, default: False
         If ``True`` additional information will be printed.
 
@@ -191,8 +199,7 @@ def build_discretized(tb_hamiltonian, discrete_coordinates,
         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)
+        tb_hamiltonian[k] = sympify(v, substitutions)
 
     discrete_coordinates = sorted(discrete_coordinates)
 

kwant/continuum/tests/test_common.py

@@ -11,28 +11,59 @@ 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
+com_A, com_B, com_C = sympy.symbols('A B C')
+x_op, y_op, z_op = 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']
+@pytest.mark.parametrize('input_expr, output_expr', [
+    ('k_x * A(x) * k_x', kx * com_A(x_op) * kx),
+    ('[[k_x * A(x) * k_x]]', sympy.Matrix([kx * com_A(x_op) * kx])),
+    ('k_x * sigma_y + k_y * sigma_x', kx * msigma(2) + ky * msigma(1)),
+    ('[[k_x*A(x)*k_x, B(x, y)*k_x], [k_x*B(x, y), C*k_y**2]]',
+            sympy.Matrix([[kx*com_A(x_op)*kx, com_B(x_op, y_op)*kx],
+                          [kx*com_B(x_op, y_op), com_C*ky**2]]))
+])
+def test_sympify(input_expr, output_expr):
+    assert sympify(input_expr) == output_expr
+    assert sympify(sympify(input_expr)) == output_expr
+
+
+
+
+@pytest.mark.parametrize('input_expr, output_expr, subs', [
+    ('k_x', kx + ky, {'k_x': 'k_x + k_y'}),
+    ('x', x_op + y_op, {'x': 'x + y'}),
+    ('A', com_A + com_B, {'A': 'A + B'}),
+    ('A', com_A + com_B(x_op), {'A': 'A + B(x)'}),
+    ('A', msigma(2), {'A': "[[0, -1j], [1j, 0]]"}),
+])
+def test_sympify_substitutions(input_expr, output_expr, subs):
+    assert sympify(input_expr, substitutions=subs) == output_expr
+    assert sympify(sympify(input_expr), substitutions=subs) == output_expr
+
+    subs = {k: sympify(v) for k, v in subs.items()}
+    assert sympify(input_expr, substitutions=subs) == output_expr
+    assert sympify(sympify(input_expr), substitutions=subs) == output_expr
 
-    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]])
+    subs = {sympify(k): sympify(v) for k, v in subs.items()}
+    assert sympify(sympify(input_expr), substitutions=subs) == output_expr
+
+
+
+
+@pytest.mark.parametrize('input_expr, output_expr, subs', [
+    ('A + k_x**2 * eye(2)',
+        kx**2 * sympy.eye(2) + msigma(2),
+        {'A': "[[0, -1j], [1j, 0]]"})
+])
+def test_sympify_mix_symbol_and_matrx(input_expr, output_expr, subs):
+    assert sympify(input_expr, substitutions=subs) == output_expr
+
+    subs = {k: sympify(v) for k, v in subs.items()}
+    assert sympify(input_expr, substitutions=subs) == output_expr
 
 
 
@@ -114,7 +145,16 @@ def test_lambdify(e, should_be, kwargs):
     assert e(**kwargs) == should_be(**kwargs)
 
 
+@pytest.mark.parametrize("e, kwargs", [
+    ("x + y", dict(x=1, y=3, z=5)),
+    (sympify("x+y"), dict(x=1, y=3, z=5))
+    ])
+def test_lambdify_substitutions(e, kwargs):
+    should_be = lambda x, y, z: x + y + z
+    subs = {'y': 'y + z'}
 
+    e = lambdify(e, substitutions=subs)
+    assert e(**kwargs) == should_be(**kwargs)
 
 
 # dispersion_string = ('A * k_x**2 * eye(2) + B * k_y**2 * eye(2)'

kwant/continuum/tests/test_discretizer.py

 import sympy
 import numpy as np
 
+from .._common import sympify
 from ..discretizer import discretize
 from ..discretizer import discretize_symbolic
 from ..discretizer import build_discretized
@@ -128,6 +129,24 @@ def test_simple_derivations(commutative):
         assert got == out
 
 
+@pytest.mark.parametrize('e_to_subs, e, subs', [
+    ('k_x', 'k_x + k_y', {'k_x': 'k_x + k_y'}),
+    ('k_x**2 + V', 'k_x**2 + V + V_0', {'V': 'V + V_0'}),
+    ('k_x**2 + A + C', 'k_x**2 + B + 5', {'A': 'B + 5', 'C': 0}),
+    ('x + y + z', '1 + 3 + 5', {'x': 1, 'y': 3, 'z': 5})
+    ])
+def test_simple_derivations_with_subs(e_to_subs, e, subs):
+    # check with strings
+    one = discretize_symbolic(e_to_subs, 'xyz', substitutions=subs)
+    two = discretize_symbolic(e, 'xyz')
+    assert one == two
+
+    # check with sympy objects
+    one = discretize_symbolic(sympify(e_to_subs), 'xyz', substitutions=subs)
+    two = discretize_symbolic(sympify(e), 'xyz')
+    assert one == two
+
+
 def test_simple_derivations_matrix():
     test = {
         kx**2                   : {(0,): 2/a**2, (1,): -1/a**2},