# Copyright 2011-2017 Kwant authors. # # This file is part of Kwant. It is subject to the license terms in the file # LICENSE.rst found in the top-level directory of this distribution and at # http://kwant-project.org/license. A list of Kwant authors can be found in # the file AUTHORS.rst at the top-level directory of this distribution and at # http://kwant-project.org/authors. import keyword from collections import defaultdict import numpy as np import sympy import sympy.abc import sympy.physics.quantum from sympy.core.function import AppliedUndef from sympy.core.sympify import converter from sympy.core.core import all_classes as sympy_classes from sympy.physics.matrices import msigma as _msigma import warnings from .._common import reraise_warnings # TODO: remove when sympy correctly includes MutableDenseMatrix (lol). sympy_classes = set(sympy_classes) | {sympy.MutableDenseMatrix} momentum_operators = sympy.symbols('k_x k_y k_z', commutative=False) position_operators = sympy.symbols('x y z', commutative=False) pauli = [sympy.eye(2), _msigma(1), _msigma(2), _msigma(3)] 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 extra_ns['I'] del extra_ns['pi'] ################ Helpers to handle sympy 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 locals : dict or ``None`` (default) Additional definitions for `~kwant.continuum.sympify`. Examples -------- >>> f = lambdify('a + b', locals={'b': 'b + c'}) >>> f(1, 3, 5) 9 >>> ns = {'sigma_plus': [[0, 2], [0, 0]]} >>> f = lambdify('k_x**2 * sigma_plus', ns) >>> f(0.25) array([[ 0. , 0.125], [ 0. , 0. ]]) """ with reraise_warnings(level=4): 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)] return sympy.lambdify(sorted(args), expr) def sympify(expr, locals=None): """Sympify object using special rules for Hamiltonians. If `'expr`` is already a type that SymPy understands, it will do nothing but return that value. Note that ``locals`` will not be used in this situation. 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 "gamma") are treated as symbols, * "kron" corresponds to ``sympy.physics.quantum.TensorProduct``, and "identity" to ``sympy.eye``, * "sigma_0", "sigma_x", "sigma_y", "sigma_z" are the Pauli matrices. In addition, Python list literals are interpreted as SymPy matrices. .. warning:: This function uses ``eval`` (because it calls ``sympy.sympify``), and thus should not be used on unsanitized input. Parameters ---------- expr : str or SymPy expression Expression to be converted to a SymPy object. locals : dict or ``None`` (default) Additional entries for the namespace under which `expr` is sympified. The keys must be valid Python variable names. 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 ------- result : SymPy object Examples -------- >>> 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', locals={'V': 'V_0 + V(x)'}) k_x**2 + V(x) + V_0 >>> 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', locals={'c': 'x'}) k_x*A(x)*k_x >>> sympify('k_x * A(c) * k_x', locals={'c': sympy.Symbol('x')}) A(x)*k_x**2 """ stored_value = None # if ``expr`` is already a ``sympy`` object we may terminate a code path if isinstance(expr, tuple(sympy_classes)): if locals: warnings.warn('Input expression is already SymPy object: ' '"locals" will not be used.', RuntimeWarning, stacklevel=2) # We assume that all present functions, like "sin", "cos", will be # provided by user during the final evaluation through "params". # Therefore we make sure they are defined as AppliedUndef, not built-in # sympy types. subs = {r: sympy.Function(str(r.func))(*r.args) for r in expr.atoms(sympy.Function)} return expr.subs(subs) # if ``expr`` is not a "sympy" then we proceed with sympifying process if locals is None: locals = {} 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 # Cast numpy array values in locals to sympy matrices to make sure they have # correct format locals = {k: (sympy.Matrix(v) if isinstance(v, np.ndarray) else 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=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. hamiltonian = sympy.sympify(hamiltonian) finally: if stored_value is not None: converter[list] = stored_value else: del converter[list] return 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] expr = expr.subs({s: sympy.Symbol(s.name) for s in symbols}) return expr def monomials(expr, gens=None): """Parse ``expr`` into monomials in the symbols in ``gens``. Parameters ---------- expr: sympy.Expr or sympy.Matrix Sympy expression to be parsed into monomials. gens: sequence of sympy.Symbol objects or strings (optional) Generators of monomials. If unset it will default to all symbols used in ``expr``. Returns ------- dictionary (generator: monomial) Examples -------- >>> expr = kwant.continuum.sympify("A * (x**2 + y) + B * x + C") >>> monomials(expr, gens=('x', 'y')) {1: C, x: B, x**2: A, y: A} """ if gens is None: gens = expr.atoms(sympy.Symbol) else: gens = [sympify(g) for g in gens] if not isinstance(expr, sympy.MatrixBase): return _expression_monomials(expr, gens) else: output = defaultdict(lambda: sympy.zeros(*expr.shape)) for (i, j), e in np.ndenumerate(expr): mons = _expression_monomials(e, gens) for key, val in mons.items(): output[key][i, j] += val return dict(output) def _expression_monomials(expr, gens): """Parse ``expr`` into monomials in the symbols in ``gens``. Parameters ---------- expr: sympy.Expr Sympy expr to be parsed. gens: sequence of sympy.Symbol Generators of monomials. Returns ------- dictionary (generator: monomial) """ expr = sympy.expand(expr) output = defaultdict(lambda: sympy.Integer(0)) for summand in expr.as_ordered_terms(): key = [] val = [] for factor in summand.as_ordered_factors(): symbol, exponent = factor.as_base_exp() if symbol in gens: key.append(factor) else: val.append(factor) output[sympy.Mul(*key)] += sympy.Mul(*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)