diff --git a/kwant/continuum/_common.py b/kwant/continuum/_common.py
index c8775e03074d3cf1e2f29c57c62765d4948e9d92..26e6dba3e5cdca6197be4ed1b3b00c535cbcdea8 100644
--- a/kwant/continuum/_common.py
+++ b/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
diff --git a/kwant/continuum/discretizer.py b/kwant/continuum/discretizer.py
index 2bbac79793488edd454a3e309a44fe3d9c1652e5..148303ee6254e061db08341f09dff1befdf13dce 100644
--- a/kwant/continuum/discretizer.py
+++ b/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)
diff --git a/kwant/continuum/tests/test_common.py b/kwant/continuum/tests/test_common.py
index ed4192a8d095d5edb7afe568d12f854f6d5853e4..875ca9a9d27fd6fc8d849130ffa8a6ba29e5f71f 100644
--- a/kwant/continuum/tests/test_common.py
+++ b/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)'
diff --git a/kwant/continuum/tests/test_discretizer.py b/kwant/continuum/tests/test_discretizer.py
index 2b6a69efbdfe7ed509d5e03d78f7e9683d117a93..47c8a6231b4e2dbf0958b291c861bfcadfa842b8 100644
--- a/kwant/continuum/tests/test_discretizer.py
+++ b/kwant/continuum/tests/test_discretizer.py
@@ -1,6 +1,7 @@
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},