Commit 68b80a81 authored by Christoph Groth's avatar Christoph Groth
Browse files

simplify behavior of continuum.symplify when given a SymPy expression

closes #125
parent dffdd545
......@@ -67,7 +67,7 @@
def substitutions():
@@ -173,15 +184,18 @@
sympify('k_x**2 * sz + alpha * k_x * sx', subs=subs),
sympify('k_x**2 * sz + alpha * k_x * sx', locals=subs),
)
- print(e[0] == e[1] == e[2])
......@@ -75,9 +75,9 @@
+ print(e[0] == e[1] == e[2], file=f)
subs = {'A': 'A(x) + B', 'V': 'V(x) + V_0', 'C': 5}
- print(sympify('k_x * A * k_x + V + C', subs=subs))
- print(sympify('k_x * A * k_x + V + C', locals=subs))
+ with open('discretizer_subs_2.txt', 'w') as f:
+ print(sympify('k_x * A * k_x + V + C', subs=subs), file=f)
+ print(sympify('k_x * A * k_x + V + C', locals=subs), file=f)
def main():
......
......@@ -144,7 +144,7 @@ def lattice_spacing():
def plot(ax, a=1):
#HIDDEN_BEGIN_ls_hk_cont
h_k = kwant.continuum.lambdify(hamiltonian, subs=params)
h_k = kwant.continuum.lambdify(hamiltonian, locals=params)
k_cont = np.linspace(-4, 4, 201)
e_cont = [scipy.linalg.eigvalsh(h_k(k_x=ki)) for ki in k_cont]
#HIDDEN_END_ls_hk_cont
......@@ -191,7 +191,7 @@ def substitutions():
e = (
sympify('[[k_x**2, alpha * k_x], [k_x * alpha, -k_x**2]]'),
sympify('k_x**2 * sigma_z + alpha * k_x * sigma_x'),
sympify('k_x**2 * sz + alpha * k_x * sx', subs=subs),
sympify('k_x**2 * sz + alpha * k_x * sx', locals=subs),
)
print(e[0] == e[1] == e[2])
......@@ -199,7 +199,7 @@ def substitutions():
#HIDDEN_BEGIN_subs_2
subs = {'A': 'A(x) + B', 'V': 'V(x) + V_0', 'C': 5}
print(sympify('k_x * A * k_x + V + C', subs=subs))
print(sympify('k_x * A * k_x + V + C', locals=subs))
#HIDDEN_END_subs_2
......
......@@ -285,7 +285,7 @@ square ``[]`` brackets. For example, all following expressions are equivalent:
.. literalinclude:: ../images/discretizer_subs_1.txt
We can use the ``substitutions`` keyword parameter to substitute expressions
We can use the ``locals`` keyword parameter to substitute expressions
and numerical values:
.. literalinclude:: discretize.py
......
......@@ -7,7 +7,7 @@
# http://kwant-project.org/authors.
import functools
import inspect
import keyword
from collections import defaultdict
from operator import mul
......@@ -26,48 +26,52 @@ position_operators = sympy.symbols('x y z', commutative=False)
pauli = [sympy.eye(2), _msigma(1), _msigma(2), _msigma(3)]
default_subs = sympy.abc._clash.copy()
default_subs.update({s.name: s for s in momentum_operators})
default_subs.update({s.name: s for s in position_operators})
default_subs.update({'kron': sympy.physics.quantum.TensorProduct,
'eye': sympy.eye, 'identity': sympy.eye})
default_subs.update({'sigma_{}'.format(c): p for c, p in zip('0xyz', pauli)})
extra_ns = sympy.abc._clash.copy()
extra_ns.update({s.name: s for s in momentum_operators})
extra_ns.update({s.name: s for s in position_operators})
extra_ns.update({'kron': sympy.physics.quantum.TensorProduct,
'eye': sympy.eye, 'identity': sympy.eye})
extra_ns.update({'sigma_{}'.format(c): p for c, p in zip('0xyz', pauli)})
# workaroud for https://github.com/sympy/sympy/issues/12060
del default_subs['I']
del default_subs['pi']
del extra_ns['I']
del extra_ns['pi']
################ Helpers to handle sympy
def lambdify(expr, subs=None):
def lambdify(expr, locals=None):
"""Return a callable object for computing a continuum Hamiltonian.
.. warning::
This function uses ``eval`` (because it calls ``sympy.sympify``), and
thus should not be used on unsanitized input.
If necessary, the given expression is sympified using
`kwant.continuum.sympify`. It is then converted into a callable object.
Parameters
----------
expr : str or SymPy expression
Expression to be converted into a callable object
subs : dict or ``None`` (default)
Namespace of substitutions to be performed on `expr`.
locals : dict or ``None`` (default)
Additional definitions for `~kwant.continuum.sympify`.
Example:
--------
>>> f = lambdify('a + b', subs={'b': 'b + c'})
>>> f = lambdify('a + b', locals={'b': 'b + c'})
>>> f(1, 3, 5)
9
>>> subs = {'sigma_plus': [[0, 2], [0, 0]]}
>>> f = lambdify('k_x**2 * sigma_plus', subs)
>>> ns = {'sigma_plus': [[0, 2], [0, 0]]}
>>> f = lambdify('k_x**2 * sigma_plus', ns)
>>> f(0.25)
array([[ 0. , 0.125],
[ 0. , 0. ]])
"""
expr = sympify(expr, subs)
expr = sympify(expr, locals)
args = [s.name for s in expr.atoms(sympy.Symbol)]
args += [str(f.func) for f in expr.atoms(AppliedUndef, sympy.Function)]
......@@ -75,17 +79,17 @@ def lambdify(expr, subs=None):
return sympy.lambdify(sorted(args), expr)
def sympify(expr, subs=None):
def sympify(expr, locals=None):
"""Sympify object using special rules for Hamiltonians.
If ``expr`` is a string, all keys in ``subs`` must be strings as well. In a
first step, all values of ``subs`` are sympified. Then, `subs` is used as
the ``locals`` argument in an internal call to ``sympy.sympify``. The
``locals`` namespace is pre-populated such that
If ``expr`` is a SymPy object, it is returned unmodified.
Otherwise, it is sympified by ``sympy.sympify`` with a modified namespace
such that
* the position operators "x", "y" or "z" and momentum operators "k_x",
"k_y", and "k_z" do not commute,
* all single-letter identifiers and names of greek letters (e.g. "pi" or
* all single-letter identifiers and names of Greek letters (e.g. "pi" or
"gamma") are treated as symbols,
* "kron" corresponds to ``sympy.physics.quantum.TensorProduct``, and
"identity" to ``sympy.eye``,
......@@ -97,27 +101,22 @@ def sympify(expr, subs=None):
This function uses ``eval`` (because it calls ``sympy.sympify``), and
thus should not be used on unsanitized input.
If `expr` is already a SymPy expression or a SymPy matrix, both the keys
and the values of `subs` are sympified (with the above special rules in
force) and used to to perform substitution using the ``.subs`` method of
`expr`.
.. note::
Any (part of) argument to this function that gets sympified (for
example the values of `subs`) may be already a SymPy object. In this
case it is left unchanged. In particular, the commutativity of its
terms is not altered. This possibly confusing effect is demonstrated
in the last example below.
Parameters
----------
expr : str or SymPy expression
Expression to be converted to a SymPy object.
subs : dict or ``None`` (default)
Namespace of substitutions to be performed on the input `expr`.
If `expr` is a string, the keys must be strings as well. Otherwise
they may be any sympifyable object. The values must be sympifyable
objects.
locals : dict or ``None`` (default)
Additional entries for the namespace under which `expr` is sympified.
The keys must be valid Python identifiers and may not be Python
keywords. The values may be strings, since they are all are sent
through `continuum.sympify` themselves before use. (Note that this is
a difference to how ``sympy.sympify`` behaves.)
.. note::
When a value of `locals` is already a SymPy object, it is used
as-is, and the caller is responsible to set the commutativity of
its symbols appropriately. This possible source of errors is
demonstrated in the last example below.
Returns
-------
......@@ -128,46 +127,44 @@ def sympify(expr, subs=None):
>>> sympify('k_x * A(x) * k_x + V(x)')
k_x*A(x)*k_x + V(x) # as valid sympy object
>>> sympify('k_x**2 + V', subs={'V': 'V_0 + V(x)'})
>>> sympify('k_x**2 + V', locals={'V': 'V_0 + V(x)'})
k_x**2 + V(x) + V_0
>>> subs = {'sigma_plus': [[0, 2], [0, 0]]}
>>> sympify('k_x**2 * sigma_plus', subs)
>>> ns = {'sigma_plus': [[0, 2], [0, 0]]}
>>> sympify('k_x**2 * sigma_plus', ns)
Matrix([
[0, 2*k_x**2],
[0, 0]])
>>> sympify('k_x * A(c) * k_x', subs={'c': 'x'})
>>> sympify('k_x * A(c) * k_x', locals={'c': 'x'})
k_x*A(x)*k_x
>>> sympify('k_x * A(c) * k_x', subs={'c': sympy.Symbol('x')})
>>> sympify('k_x * A(c) * k_x', locals={'c': sympy.Symbol('x')})
A(x)*k_x**2
"""
stored_value = None
sympified_types = (sympy.Expr, sympy.matrices.MatrixBase)
if subs is None:
subs = {}
if locals is None:
locals = {}
# if ``expr`` is already a ``sympy`` object we make use of ``subs``
# and terminate a code path.
if isinstance(expr, sympified_types):
subs = {sympify(k): sympify(v) for k, v in subs.items()}
return expr.subs(subs)
# if ``expr`` is not a sympified type then we proceed with sympifying
# process, we expect all keys in ``subs`` to be strings at this moment.
if not all(isinstance(k, str) for k in subs):
raise ValueError("If 'expr' is not already a sympy object ",
"then keys of 'subs' must be strings.")
# sympify values of subs before updating it with default_subs
subs = {k: sympify(v) for k, v in subs.items()}
for k, v in default_subs.items():
subs.setdefault(k, v)
return expr
for k in locals:
if (not isinstance(k, str)
or not k.isidentifier() or keyword.iskeyword(k)):
raise ValueError(
"Invalid key in 'locals': {}\nKeys must be "
"identifiers and may not be keywords".format(repr(k)))
# sympify values of locals before updating it with extra_ns
locals = {k: sympify(v) for k, v in locals.items()}
for k, v in extra_ns.items():
locals.setdefault(k, v)
try:
stored_value = converter.pop(list, None)
converter[list] = lambda x: sympy.Matrix(x)
hamiltonian = sympy.sympify(expr, locals=subs)
hamiltonian = sympy.sympify(expr, locals=locals)
# 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.
......@@ -195,8 +192,7 @@ def make_commutative(expr, *symbols):
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)
expr = expr.subs({s: sympy.Symbol(s.name) for s in symbols})
return expr
......
......@@ -81,9 +81,14 @@ class _DiscretizedBuilder(builder.Builder):
################ Interface functions
def discretize(hamiltonian, coords=None, *, grid_spacing=1,
subs=None):
locals=None):
"""Construct a tight-binding model from a continuum Hamiltonian.
If necessary, the given Hamiltonian is sympified using
`kwant.continuum.sympify`. It is then discretized symbolically and turned
into a `~kwant.builder.Builder` instance that may be used with
`~kwant.builder.Builder.fill`.
This is a convenience function that is equivalent to first calling
`~kwant.continuum.discretize_symbolic` and feeding its result into
`~kwant.continuum.build_discretized`.
......@@ -94,8 +99,8 @@ def discretize(hamiltonian, coords=None, *, grid_spacing=1,
Parameters
----------
hamiltonian : string, or sympy.Expr, or sympy.Matrix
Symbolic representation of a continous Hamiltonian. It is
hamiltonian : str, or sympy.Expr, or sympy.Matrix
Symbolic representation of a continuous Hamiltonian. It is
converted to a SymPy expression using `kwant.continuum.sympify`.
coords : sequence of strings, or ``None`` (default)
The coordinates for which momentum operators will be treated as
......@@ -104,31 +109,33 @@ def discretize(hamiltonian, coords=None, *, grid_spacing=1,
Hamiltonian by reading the present coordinates and momentum operators.
grid_spacing : int or float, default: 1
Spacing of the (quadratic or cubic) discretization grid.
subs : dict or ``None`` (default)
A namespace of substitutions to be performed on `hamiltonian`. See
`kwant.continuum.sympify` for details. May be used to simplify input of
matrices or alternate input before proceeding further. For example:
``subs={'k': 'k_x + I * k_y'}`` or
``subs={'sigma_plus': [[0, 2], [0, 0]]}``.
locals : dict or ``None`` (default)
Additional namespace entries for `~kwant.continuum.sympify`. May be
used to simplify input of matrices or modify input before proceeding
further. For example:
``locals={'k': 'k_x + I * k_y'}`` or
``locals={'sigma_plus': [[0, 2], [0, 0]]}``.
Returns
-------
model: `~kwant.builder.Builder`
The translationally symmetric builder that corresponds to the provided
Hamiltonian. It may be passed directly into the method
`~kwant.builder.Builder.fill` of a builder to construct a system of a
desired shape.
Hamiltonian.
"""
tb, coords = discretize_symbolic(hamiltonian, coords, subs=subs)
tb, coords = discretize_symbolic(hamiltonian, coords, locals=locals)
return build_discretized(tb, coords, grid_spacing=grid_spacing)
def discretize_symbolic(hamiltonian, coords=None, *, subs=None):
def discretize_symbolic(hamiltonian, coords=None, *, locals=None):
"""Discretize a continuous Hamiltonian into a tight-binding representation.
The two objects returned by this function may be used directly as the first
two arguments for `~kwant.continuum.build_discretized`.
If necessary, the given Hamiltonian is sympified using
`kwant.continuum.sympify`. It is then discretized symbolically.
The two return values may be used directly as the first two arguments for
`~kwant.continuum.build_discretized`.
.. warning::
This function uses ``eval`` (because it calls ``sympy.sympify``), and
......@@ -136,20 +143,20 @@ def discretize_symbolic(hamiltonian, coords=None, *, subs=None):
Parameters
----------
hamiltonian : string, or sympy.Expr, or sympy.Matrix
Symbolic representation of a continous Hamiltonian. It is
converted to a sympy expression using `kwant.continuum.sympify`.
hamiltonian : str, or sympy.Expr, or sympy.Matrix
Symbolic representation of a continuous Hamiltonian. It is
converted to a SymPy expression using `kwant.continuum.sympify`.
coords : sequence of strings, or ``None`` (default)
The coordinates for which momentum operators will be treated as
differential operators. May contain only "x", "y" and "z" and must be
sorted. If not provided, `coords` will be obtained from the input
Hamiltonian by reading the present coordinates and momentum operators.
subs : dict, defaults to empty
A namespace of substitutions to be performed on `hamiltonian`. See
`kwant.continuum.sympify` for details. May be used to simplify input of
matrices or alternate input before proceeding further. For example:
``subs={'k': 'k_x + I * k_y'}`` or
``subs={'sigma_plus': [[0, 2], [0, 0]]}``.
locals : dict or ``None`` (default)
Additional namespace entries for `~kwant.continuum.sympify`. May be
used to simplify input of matrices or modify input before proceeding
further. For example:
``locals={'k': 'k_x + I * k_y'}`` or
``locals={'sigma_plus': [[0, 2], [0, 0]]}``.
Returns
-------
......@@ -160,7 +167,7 @@ def discretize_symbolic(hamiltonian, coords=None, *, subs=None):
The coordinates that have been discretized.
"""
hamiltonian = sympify(hamiltonian, subs)
hamiltonian = sympify(hamiltonian, locals)
atoms_names = [s.name for s in hamiltonian.atoms(sympy.Symbol)]
if any( s == 'a' for s in atoms_names):
......@@ -215,9 +222,14 @@ def discretize_symbolic(hamiltonian, coords=None, *, subs=None):
return tb, coords
def build_discretized(tb_hamiltonian, coords, *, grid_spacing=1, subs=None):
def build_discretized(tb_hamiltonian, coords, *, grid_spacing=1, locals=None):
"""Create a template builder from a symbolic tight-binding Hamiltonian.
The provided symbolic tight-binding Hamiltonian is put on a (hyper) square
lattice and turned into Python functions. These functions are used to
create a `~kwant.builder.Builder` instance that may be used with
`~kwant.builder.Builder.fill` to construct a system of a desired shape.
The return values of `~kwant.continuum.discretize_symbolic` may be used
directly for the first two arguments of this function.
......@@ -238,27 +250,26 @@ def build_discretized(tb_hamiltonian, coords, *, grid_spacing=1, subs=None):
sorted.
grid_spacing : int or float, default: 1
Spacing of the (quadratic or cubic) discretization grid.
subs : dict, defaults to empty
A namespace of substitutions to be performed on the values of
`tb_hamiltonian`. See `kwant.continuum.sympify` for details. May be
used to simplify input of matrices or alternate input before proceeding
further. For example: ``subs={'k': 'k_x + I * k_y'}`` or
``subs={'sigma_plus': [[0, 2], [0, 0]]}``.
locals : dict, defaults to empty
Additional namespace entries for the calls of
`~kwant.continuum.sympify` on the values of `tb_hamiltonian`. May be
used to simplify input of matrices or modify input before proceeding
further. For example:
``locals={'k': 'k_x + I * k_y'}`` or
``locals={'sigma_plus': [[0, 2], [0, 0]]}``.
Returns
-------
model : `~kwant.builder.Builder`
The translationally symmetric builder that corresponds to the provided
Hamiltonian. It may be passed directly into the method
`~kwant.builder.Builder.fill` of a builder to construct a system of a
desired shape.
Hamiltonian.
"""
if len(coords) == 0:
raise ValueError('Discrete coordinates cannot be empty.')
for k, v in tb_hamiltonian.items():
tb_hamiltonian[k] = sympify(v, subs)
tb_hamiltonian[k] = sympify(v, locals)
coords = list(coords)
if coords != sorted(coords):
......
......@@ -53,15 +53,21 @@ def test_sympify(input_expr, output_expr):
('A', msigma(2), {'A': "[[0, -1j], [1j, 0]]"}),
])
def test_sympify_substitutions(input_expr, output_expr, subs):
assert sympify(input_expr, subs=subs) == output_expr
assert sympify(sympify(input_expr), subs=subs) == output_expr
assert sympify(input_expr, locals=subs) == output_expr
subs = {k: sympify(v) for k, v in subs.items()}
assert sympify(input_expr, subs=subs) == output_expr
assert sympify(sympify(input_expr), subs=subs) == output_expr
assert sympify(input_expr, locals=subs) == output_expr
subs = {sympify(k): sympify(v) for k, v in subs.items()}
assert sympify(sympify(input_expr), subs=subs) == output_expr
@pytest.mark.parametrize('subs', [
{1: 2},
{'1': 2},
{'for': 33},
{'x': 1, '1j': 7}
])
def test_sympify_invalid_substitutions(subs):
with pytest.raises(ValueError):
sympify('x + y', locals=subs)
@pytest.mark.parametrize('input_expr, output_expr, subs', [
......@@ -74,10 +80,10 @@ def test_sympify_substitutions(input_expr, output_expr, subs):
])
def test_sympify_mix_symbol_and_matrx(input_expr, output_expr, subs):
assert sympify(input_expr, subs=subs) == output_expr
assert sympify(input_expr, locals=subs) == output_expr
subs = {k: sympify(v) for k, v in subs.items()}
assert sympify(input_expr, subs=subs) == output_expr
assert sympify(input_expr, locals=subs) == output_expr
A, B, non_x = sympy.symbols('A B x', commutative=False)
......@@ -173,11 +179,11 @@ def test_lambdify(e, 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)),
(sympify("x+y+z"), 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, subs=subs)
e = lambdify(e, locals=subs)
assert e(**kwargs) == should_be(**kwargs)
......@@ -14,7 +14,6 @@ import pytest
import sympy
from .._common import sympify
from ..discretizer import discretize
from ..discretizer import discretize_symbolic
from ..discretizer import build_discretized
......@@ -134,7 +133,7 @@ def test_simple_derivations(commutative):
assert got == out
for inp, out in test.items():
got, _ = discretize_symbolic(str(inp), subs=ns)
got, _ = discretize_symbolic(str(inp), locals=ns)
assert got == out
......@@ -146,16 +145,10 @@ def test_simple_derivations(commutative):
('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', subs=subs)
one = discretize_symbolic(e_to_subs, 'xyz', locals=subs)
two = discretize_symbolic(e, 'xyz')
assert one == two
# check with sympy objects
one = discretize_symbolic(sympify(e_to_subs), 'xyz', subs=subs)
two = discretize_symbolic(sympify(e), 'xyz')
assert one == two
def test_simple_derivations_matrix():
test = {
......@@ -191,11 +184,11 @@ def test_simple_derivations_matrix():
assert got == out
for inp, out in new_test:
got, _ = discretize_symbolic(str(inp), subs=ns)
got, _ = discretize_symbolic(str(inp), locals=ns)
assert got == out
for inp, out in new_test:
got, _ = discretize_symbolic(str(inp).replace('Matrix', ''), subs=ns)
got, _ = discretize_symbolic(str(inp).replace('Matrix', ''), locals=ns)
assert got == out
......@@ -400,7 +393,7 @@ def test_numeric_functions_basic_string():
def test_numeric_functions_with_subs(e_to_subs, e, subs):
p = {'A': 1, 'B': 2}
builder_direct = discretize(e)
builder_subs = discretize(e_to_subs, subs=subs)
builder_subs = discretize(e_to_subs, locals=subs)
lat = next(iter(builder_direct.sites()))[0]
assert builder_direct[lat(0)](None, **p) == builder_subs[lat(0)](None, **p)
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment