Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright 2011-2017 Kwant authors.
2#
3# This file is part of Kwant. It is subject to the license terms in the file
4# LICENSE.rst found in the top-level directory of this distribution and at
5# https://kwant-project.org/license. A list of Kwant authors can be found in
6# the file AUTHORS.rst at the top-level directory of this distribution and at
7# https://kwant-project.org/authors.
9import keyword
10from collections import defaultdict
12import numpy as np
14import sympy
15import sympy.abc
16import sympy.physics.quantum
17from sympy.core.function import AppliedUndef
18from sympy.core.sympify import converter
19from sympy.core.core import all_classes as sympy_classes
20from sympy.physics.matrices import msigma as _msigma
22import warnings
24from .._common import reraise_warnings
26# TODO: remove when sympy correctly includes MutableDenseMatrix (lol).
27sympy_classes = set(sympy_classes) | {sympy.MutableDenseMatrix}
29momentum_operators = sympy.symbols('k_x k_y k_z', commutative=False)
30position_operators = sympy.symbols('x y z', commutative=False)
32pauli = [sympy.eye(2), _msigma(1), _msigma(2), _msigma(3)]
34extra_ns = sympy.abc._clash.copy()
35extra_ns.update({s.name: s for s in momentum_operators})
36extra_ns.update({s.name: s for s in position_operators})
37extra_ns.update({'kron': sympy.physics.quantum.TensorProduct,
38 'eye': sympy.eye, 'identity': sympy.eye})
39extra_ns.update({'sigma_{}'.format(c): p for c, p in zip('0xyz', pauli)})
42# TODO: remove this once https://github.com/sympy/sympy/issues/12060 is fixed
43del extra_ns['I']
44del extra_ns['pi']
47################ Helpers to handle sympy
49def lambdify(expr, locals=None):
50 """Return a callable object for computing a continuum Hamiltonian.
52 .. warning::
53 This function uses ``eval`` (because it calls ``sympy.sympify``), and
54 thus should not be used on unsanitized input.
56 If necessary, the given expression is sympified using
57 `kwant.continuum.sympify`. It is then converted into a callable object.
59 Parameters
60 ----------
61 expr : str or SymPy expression
62 Expression to be converted into a callable object
63 locals : dict or ``None`` (default)
64 Additional definitions for `~kwant.continuum.sympify`.
66 Examples
67 --------
68 >>> f = lambdify('a + b', locals={'b': 'b + c'})
69 >>> f(1, 3, 5)
70 9
72 >>> ns = {'sigma_plus': [[0, 2], [0, 0]]}
73 >>> f = lambdify('k_x**2 * sigma_plus', ns)
74 >>> f(0.25)
75 array([[ 0. , 0.125],
76 [ 0. , 0. ]])
77 """
78 with reraise_warnings(level=4):
79 expr = sympify(expr, locals)
81 args = [s.name for s in expr.atoms(sympy.Symbol)]
82 args += [str(f.func) for f in expr.atoms(AppliedUndef, sympy.Function)]
84 return sympy.lambdify(sorted(args), expr)
87def sympify(expr, locals=None):
88 """Sympify object using special rules for Hamiltonians.
90 If `'expr`` is already a type that SymPy understands, it will do nothing
91 but return that value. Note that ``locals`` will not be used in this
92 situation.
94 Otherwise, it is sympified by ``sympy.sympify`` with a modified namespace
95 such that
97 * the position operators "x", "y" or "z" and momentum operators "k_x",
98 "k_y", and "k_z" do not commute,
99 * all single-letter identifiers and names of Greek letters (e.g. "pi" or
100 "gamma") are treated as symbols,
101 * "kron" corresponds to ``sympy.physics.quantum.TensorProduct``, and
102 "identity" to ``sympy.eye``,
103 * "sigma_0", "sigma_x", "sigma_y", "sigma_z" are the Pauli matrices.
105 In addition, Python list literals are interpreted as SymPy matrices.
107 .. warning::
108 This function uses ``eval`` (because it calls ``sympy.sympify``), and
109 thus should not be used on unsanitized input.
111 Parameters
112 ----------
113 expr : str or SymPy expression
114 Expression to be converted to a SymPy object.
115 locals : dict or ``None`` (default)
116 Additional entries for the namespace under which `expr` is sympified.
117 The keys must be valid Python variable names. The values may be
118 strings, since they are all are sent through `continuum.sympify`
119 themselves before use. (Note that this is a difference to how
120 ``sympy.sympify`` behaves.)
122 .. note::
123 When a value of `locals` is already a SymPy object, it is used
124 as-is, and the caller is responsible to set the commutativity of
125 its symbols appropriately. This possible source of errors is
126 demonstrated in the last example below.
128 Returns
129 -------
130 result : SymPy object
132 Examples
133 --------
134 >>> sympify('k_x * A(x) * k_x + V(x)')
135 k_x*A(x)*k_x + V(x) # as valid sympy object
137 >>> sympify('k_x**2 + V', locals={'V': 'V_0 + V(x)'})
138 k_x**2 + V(x) + V_0
140 >>> ns = {'sigma_plus': [[0, 2], [0, 0]]}
141 >>> sympify('k_x**2 * sigma_plus', ns)
142 Matrix([
143 [0, 2*k_x**2],
144 [0, 0]])
146 >>> sympify('k_x * A(c) * k_x', locals={'c': 'x'})
147 k_x*A(x)*k_x
148 >>> sympify('k_x * A(c) * k_x', locals={'c': sympy.Symbol('x')})
149 A(x)*k_x**2
151 """
152 stored_value = None
154 # if ``expr`` is already a ``sympy`` object we may terminate a code path
155 if isinstance(expr, tuple(sympy_classes)):
156 if locals:
157 warnings.warn('Input expression is already SymPy object: '
158 '"locals" will not be used.',
159 RuntimeWarning,
160 stacklevel=2)
162 # We assume that all present functions, like "sin", "cos", will be
163 # provided by user during the final evaluation through "params".
164 # Therefore we make sure they are defined as AppliedUndef, not built-in
165 # sympy types.
166 subs = {r: sympy.Function(str(r.func))(*r.args)
167 for r in expr.atoms(sympy.Function)}
169 return expr.subs(subs)
171 # if ``expr`` is not a "sympy" then we proceed with sympifying process
172 if locals is None:
173 locals = {}
175 for k in locals:
176 if (not isinstance(k, str)
177 or not k.isidentifier() or keyword.iskeyword(k)):
178 raise ValueError(
179 "Invalid key in 'locals': {}\nKeys must be "
180 "identifiers and may not be keywords".format(repr(k)))
182 # sympify values of locals before updating it with extra_ns
183 # Cast numpy array values in locals to sympy matrices to make sure they have
184 # correct format
185 locals = {k: (sympy.Matrix(v) if isinstance(v, np.ndarray) else sympify(v))
186 for k, v in locals.items()}
188 for k, v in extra_ns.items():
189 locals.setdefault(k, v)
190 try:
191 stored_value = converter.pop(list, None)
192 converter[list] = lambda x: sympy.Matrix(x)
193 hamiltonian = sympy.sympify(expr, locals=locals)
194 # if input is for example ``[[k_x * A(x) * k_x]]`` after the first
195 # sympify we are getting list of sympy objects, so we call sympify
196 # second time to obtain ``sympy`` matrices.
197 hamiltonian = sympy.sympify(hamiltonian)
198 finally:
199 if stored_value is not None: 199 ↛ 200line 199 didn't jump to line 200, because the condition on line 199 was never true
200 converter[list] = stored_value
201 else:
202 del converter[list]
204 return hamiltonian
207def make_commutative(expr, *symbols):
208 """Make sure that specified symbols are defined as commutative.
210 Parameters
211 ----------
212 expr: sympy.Expr or sympy.Matrix
213 symbols: sequace of symbols
214 Set of symbols that are requiered to be commutative. It doesn't matter
215 of symbol is provided as commutative or not.
217 Returns
218 -------
219 input expression with all specified symbols changed to commutative.
220 """
221 symbols = [sympy.Symbol(s.name, commutative=False) for s in symbols]
222 expr = expr.subs({s: sympy.Symbol(s.name) for s in symbols})
223 return expr
226def monomials(expr, gens=None):
227 """Parse ``expr`` into monomials in the symbols in ``gens``.
229 Parameters
230 ----------
231 expr: sympy.Expr or sympy.Matrix
232 Sympy expression to be parsed into monomials.
233 gens: sequence of sympy.Symbol objects or strings (optional)
234 Generators of monomials. If unset it will default to all
235 symbols used in ``expr``.
237 Returns
238 -------
239 dictionary (generator: monomial)
241 Examples
242 --------
243 >>> expr = kwant.continuum.sympify("A * (x**2 + y) + B * x + C")
244 >>> monomials(expr, gens=('x', 'y'))
245 {1: C, x: B, x**2: A, y: A}
246 """
247 if gens is None:
248 gens = expr.atoms(sympy.Symbol)
249 else:
250 gens = [sympify(g) for g in gens]
252 if not isinstance(expr, sympy.MatrixBase):
253 return _expression_monomials(expr, gens)
254 else:
255 output = defaultdict(lambda: sympy.zeros(*expr.shape))
256 for (i, j), e in np.ndenumerate(expr):
257 mons = _expression_monomials(e, gens)
258 for key, val in mons.items():
259 output[key][i, j] += val
260 return dict(output)
263def _expression_monomials(expr, gens):
264 """Parse ``expr`` into monomials in the symbols in ``gens``.
266 Parameters
267 ----------
268 expr: sympy.Expr
269 Sympy expr to be parsed.
270 gens: sequence of sympy.Symbol
271 Generators of monomials.
273 Returns
274 -------
275 dictionary (generator: monomial)
276 """
277 expr = sympy.expand(expr)
278 output = defaultdict(lambda: sympy.Integer(0))
279 for summand in expr.as_ordered_terms():
280 key = []
281 val = []
282 for factor in summand.as_ordered_factors():
283 symbol, exponent = factor.as_base_exp()
284 if symbol in gens:
285 key.append(factor)
286 else:
287 val.append(factor)
288 output[sympy.Mul(*key)] += sympy.Mul(*val)
290 return dict(output)
293################ general help functions
295def gcd(*args):
296 if len(args) == 1:
297 return args[0]
299 L = list(args)
301 while len(L) > 1:
302 a = L[len(L) - 2]
303 b = L[len(L) - 1]
304 L = L[:len(L) - 2]
306 while a:
307 a, b = b%a, a
309 L.append(b)
311 return abs(b)