Skip to content
Snippets Groups Projects
Commit 6009151a authored by Rafal Skolasinski's avatar Rafal Skolasinski
Browse files

integrate discretizer into kwant

parent f833d2c9
No related branches found
No related tags found
No related merge requests found
Expand all Hide whitespace changes
Inline Side-by-side
Showing
with 1416 additions and 1158 deletions
.gitlab-ci.yml
+ 1
− 0
View file @ 6009151a
...@@ -62,6 +62,7 @@ build documentation: ...@@ -62,6 +62,7 @@ build documentation:
run tests: run tests:
stage: test stage: test
script: script:
- pip3 install sympy
- py.test --cov=kwant --cov-report term --cov-report html --flakes kwant - py.test --cov=kwant --cov-report term --cov-report html --flakes kwant
artifacts: artifacts:
paths: paths:
......
discretizer/__init__.py deleted 100644 → 0
+ 0
− 10
View file @ f833d2c9
__all__ = ['algorithms']
from.discretizer import Discretizer
import sympy
__version__ = '0.4.3'
momentum_operators = sympy.symbols('k_x k_y k_z', commutative=False)
coordinates = sympy.symbols('x y z', commutative=False)
discretizer/algorithms.py deleted 100644 → 0
+ 0
− 416
View file @ f833d2c9
from __future__ import print_function, division
import itertools
import sympy
import numpy as np
from collections import defaultdict
from .postprocessing import make_kwant_functions
from .postprocessing import offset_to_direction
from .interpolation import interpolate_tb_hamiltonian
try:
# normal situation
from kwant.lattice import Monatomic
except ImportError:
# probably run on gitlab-ci
pass
# ************************** Some globals *********************************
wavefunction_name = 'Psi'
# **************** Operation on sympy expressions **************************
def read_coordinates(expression):
"""Read coordinates used in expression.
This function is used if ``discrete_coordinates`` are not provided by user.
Parameters:
-----------
expression : sympy.Expr or sympy.Matrix instance
Returns:
--------
discrete_coordinates : set of strings
"""
discrete_coordinates = set()
for a in expression.atoms(sympy.Symbol):
if a.name in ['k_x', 'k_y', 'k_z']:
discrete_coordinates.add(a.name.split('_')[1])
if a.name in ['x', 'y', 'z']:
discrete_coordinates.add(a.name)
return discrete_coordinates
def split_factors(expression, discrete_coordinates):
""" Split symbolic `expression` for a discretization step.
Parameters:
-----------
expression : sympy.Expr instance
The expression to be split. It should represents single summand.
Output:
-------
lhs : sympy.Expr instance
Part of expression standing to the left from operators
that acts in current discretization step.
operators: sympy.Expr instance
Operator that perform discretization in current step.
rhs : sympy.Expr instance
Part of expression that is derivated in current step.
Raises:
-------
AssertionError
if input `expression` is of type ``sympy.Add``
"""
assert not isinstance(expression, sympy.Add), \
'Input expression must not be sympy.Add. It should be a single summand.'
momentum_names = ['k_{}'.format(s) for s in discrete_coordinates]
momentum_operators = sympy.symbols(momentum_names, commutative=False)
momentum_operators += sympy.symbols(momentum_names, commutative=True)
output = {'rhs': [1], 'operator': [1], 'lhs': [1]}
if isinstance(expression, sympy.Pow):
base, exponent = expression.args
if base in momentum_operators:
output['operator'].append(base)
output['lhs'].append(sympy.Pow(base, exponent-1))
else:
output['rhs'].append(expression)
elif isinstance(expression, (int, float, sympy.Integer, sympy.Float)):
output['rhs'].append(expression)
elif isinstance(expression, (sympy.Symbol, sympy.Function)):
if expression in momentum_operators:
output['operator'].append(expression)
else:
output['rhs'].append(expression)
elif isinstance(expression, sympy.Mul):
iterator = iter(expression.args[::-1])
for factor in iterator:
if factor in momentum_operators:
output['operator'].append(factor)
break
elif factor.func == sympy.Pow and factor.args[0] in momentum_operators:
base, exponent = factor.args
output['operator'].append(base)
output['lhs'].append(sympy.Pow(base, exponent-1))
break
else:
output['rhs'].append(factor)
for factor in iterator:
output['lhs'].append(factor)
output = tuple(sympy.Mul(*output[key][::-1])
for key in ['lhs', 'operator', 'rhs'])
return output
def derivate(expression, operator):
""" Calculate derivate of expression for given momentum operator:
Parameters:
-----------
expression : sympy.Expr instance
Sympy expression containing functions to to be derivated.
operator : sympy.Symbol
Sympy symbol representing momentum operator.
Returns:
--------
output : sympy.Expr instance
Derivated input expression.
"""
if not isinstance(operator, sympy.Symbol):
raise TypeError("Input operator '{}' is not type sympy.Symbol.")
if operator.name not in ['k_x', 'k_y', 'k_z']:
raise ValueError("Input operator '{}' unkown.".format(operator))
if isinstance(expression, (int, float, sympy.Symbol)):
return 0
else:
coordinate_name = operator.name.split('_')[1]
ct = sympy.Symbol(coordinate_name, commutative=True)
cf = sympy.Symbol(coordinate_name, commutative=False)
h = sympy.Symbol('a_'+coordinate_name)
expr1 = expression.subs({ct: ct + h, cf: cf + h})
expr2 = expression.subs({ct: ct - h, cf: cf - h})
output = (expr1 - expr2) / 2 / h
return -sympy.I * sympy.expand(output)
def _discretize_summand(summand, discrete_coordinates):
""" Discretize one summand. """
assert not isinstance(summand, sympy.Add), "Input should be one summand."
def do_stuff(expr):
""" Derivate expr recursively. """
expr = sympy.expand(expr)
if isinstance(expr, sympy.Add):
return do_stuff(expr.args[-1]) + do_stuff(sympy.Add(*expr.args[:-1]))
lhs, operator, rhs = split_factors(expr, discrete_coordinates)
if rhs == 1 and operator != 1:
return 0
elif operator == 1:
return lhs*rhs
elif lhs == 1:
return derivate(rhs, operator)
else:
return do_stuff(lhs*derivate(rhs, operator))
return do_stuff(summand)
def _discretize_expression(expression, discrete_coordinates):
""" Discretize continous `expression` into discrete tb representation.
Parameters:
-----------
expression : sympy.Expr instance
The expression to be discretized.
Returns:
--------
discrete_expression: dict
dict in which key is offset of hopping ((0, 0, 0) for onsite)
and value is corresponding symbolic hopping (onsite).
Note:
-----
Recursive derivation implemented in _discretize_summand is applied
on every summand. Shortening is applied before return on output.
"""
if isinstance(expression, (int, float, sympy.Integer, sympy.Float)):
n = len(discrete_coordinates)
return {(tuple(0 for i in range(n))): expression}
if not isinstance(expression, sympy.Expr):
raise TypeError('Input expression should be a valid sympy expression.')
coordinates_names = sorted(list(discrete_coordinates))
coordinates = [sympy.Symbol(c, commutative=False) for c in coordinates_names]
wf = sympy.Function(wavefunction_name)(*coordinates)
if wf in expression.atoms(sympy.Function):
raise ValueError("Input expression must not contain {}.".format(wf))
expression = sympy.expand(expression*wf)
# make sure we have list of summands
summands = expression.args if expression.func == sympy.Add else [expression]
# discretize every summand
outputs = []
for summand in summands:
out = _discretize_summand(summand, discrete_coordinates)
out = extract_hoppings(out, discrete_coordinates)
outputs.append(out)
# gather together
discrete_expression = defaultdict(int)
for summand in outputs:
for k, v in summand.items():
discrete_expression[k] += v
return dict(discrete_expression)
def discretize(hamiltonian, discrete_coordinates):
""" Discretize continous `expression` into discrete tb representation.
Parameters:
-----------
hamiltonian : sympy.Expr or sympy.Matrix instance
The expression for the Hamiltonian.
Returns:
--------
discrete_hamiltonian: dict
dict in which key is offset of hopping ((0, 0, 0) for onsite)
and value is corresponding symbolic hopping (onsite).
Note:
-----
Recursive derivation implemented in _discretize_summand is applied
on every summand. Shortening is applied before return on output.
"""
if not isinstance(hamiltonian, sympy.matrices.MatrixBase):
onsite_zeros = (0,)*len(discrete_coordinates)
discrete_hamiltonian = {onsite_zeros: sympy.Integer(0)}
hoppings = _discretize_expression(hamiltonian, discrete_coordinates)
discrete_hamiltonian.update(hoppings)
return discrete_hamiltonian
shape = hamiltonian.shape
discrete_hamiltonian = defaultdict(lambda: sympy.zeros(*shape))
for i,j in itertools.product(range(shape[0]), range(shape[1])):
expression = hamiltonian[i, j]
hoppings = _discretize_expression(expression, discrete_coordinates)
for offset, hop in hoppings.items():
discrete_hamiltonian[offset][i,j] += hop
return discrete_hamiltonian
# ****** extracting hoppings ***********
def read_hopping_from_wf(wf):
"""Read offset of a wave function in respect to (x,y,z).
Parameters:
----------
wf : sympy.function.AppliedUndef instance
Function representing correct wave function used in discretizer.
Should be created using global `wavefunction_name`.
Returns:
--------
offset : tuple
tuple of integers or floats that represent offset in respect to (x,y,z).
Raises:
-------
ValueError
If arguments of wf are repeated / do not stand for valid coordinates or
lattice constants / order of dimensions is not lexical.
TypeError:
If wf is not of type sympy.function.AppliedUndef or its name does not
corresponds to global 'wavefunction_name'.
"""
if not isinstance(wf, sympy.function.AppliedUndef):
raise TypeError('Input should be of type sympy.function.AppliedUndef.')
if not wf.func.__name__ == wavefunction_name:
msg = 'Input should be function that represents wavefunction in module.'
raise TypeError(msg)
# Check if input is consistent and in lexical order.
# These are more checks for internal usage.
coordinates_names = ['x', 'y', 'z']
lattice_const_names = ['a_x', 'a_y', 'a_z']
arg_coords = []
for arg in wf.args:
names = [s.name for s in arg.atoms(sympy.Symbol)]
ind = -1
for s in names:
if not any(s in coordinates_names for s in names):
raise ValueError("Wave function argument '{}' is incorrect.".format(s))
if s not in coordinates_names and s not in lattice_const_names:
raise ValueError("Wave function argument '{}' is incorrect.".format(s))
if s in lattice_const_names:
s = s.split('_')[1]
tmp = coordinates_names.index(s)
if tmp in arg_coords:
msg = "Wave function '{}' arguments are inconsistent."
raise ValueError(msg.format(wf))
if ind != -1:
if tmp != ind:
msg = "Wave function '{}' arguments are inconsistent."
raise ValueError(msg.format(wf))
else:
ind = tmp
arg_coords.append(ind)
if arg_coords != sorted(arg_coords):
msg = "Coordinates of wave function '{}' are not in lexical order."
raise ValueError(msg.format(wf))
# function real body
offset = []
for argument in wf.args:
temp = sympy.expand(argument)
if temp in sympy.symbols('x y z', commutative = False):
offset.append(0)
elif temp.func == sympy.Add:
for arg_summands in temp.args:
if arg_summands.func == sympy.Mul:
if len(arg_summands.args) > 2:
print('More than two factors in an argument of wf')
if not arg_summands.args[0] in sympy.symbols('a_x a_y a_z'):
offset.append(arg_summands.args[0])
else:
offset.append(arg_summands.args[1])
elif arg_summands in sympy.symbols('a_x a_y a_z'):
offset.append(1)
else:
print('Argument of \wf is neither a sum nor a single space variable.')
return tuple(offset)
def extract_hoppings(expression, discrete_coordinates):
"""Extract hopping and perform shortening operation. """
# make sure we have list of summands
expression = sympy.expand(expression)
summands = expression.args if expression.func == sympy.Add else [expression]
hoppings = defaultdict(int)
for summand in summands:
if summand.func.__name__ == wavefunction_name:
hoppings[read_hopping_from_wf(summand)] += 1
else:
for i in range(len(summand.args)):
if summand.args[i].func.__name__ == wavefunction_name:
index = i
if index < len(summand.args) - 1:
print('Psi is not in the very end of the term. Output will be wrong!')
hoppings[read_hopping_from_wf(summand.args[-1])] += sympy.Mul(*summand.args[:-1])
# START: shortenig
discrete_coordinates = sorted(list(discrete_coordinates))
tmps = ['a_{}'.format(s) for s in discrete_coordinates]
lattice_constants = sympy.symbols(tmps)
a = sympy.Symbol('a')
# make a list of all hopping kinds we have to consider during the shortening
hops_kinds = np.array(list(hoppings))
# find the longest hopping range in each direction
longest_ranges = [np.max(hops_kinds[:,i]) for i in range(len(hops_kinds[0,:]))]
# define an array in which we are going to store by which factor we
# can shorten the hoppings in each direction
shortening_factors = np.ones_like(longest_ranges)
# Loop over the direction and each potential shortening factor.
# Inside the loop test whether the hopping distances are actually
# multiples of the potential shortening factor.
for dim in np.arange(len(longest_ranges)):
for factor in np.arange(longest_ranges[dim])+1:
modulos = np.mod(hops_kinds[:, dim], factor)
if np.sum(modulos) < 0.1:
shortening_factors[dim] = factor
# Apply the shortening factors on the hopping.
short_hopping = {}
for hopping_kind in hoppings.keys():
short_hopping_kind = tuple(np.array(hopping_kind) / shortening_factors)
for i in short_hopping_kind:
if isinstance(i, float):
assert i.is_integer()
short_hopping_kind = tuple(int(i) for i in short_hopping_kind)
short_hopping[short_hopping_kind] = hoppings[hopping_kind]
for lat_const, factor in zip(lattice_constants, shortening_factors):
factor = int(factor)
subs = {lat_const: lat_const/factor}
short_hopping[short_hopping_kind] = short_hopping[short_hopping_kind].subs(subs)
# We don't need separate a_x, a_y and a_z anymore.
for key, val in short_hopping.items():
short_hopping[key] = val.subs({i: a for i in sympy.symbols('a_x a_y a_z')})
return short_hopping
discretizer/discretizer.py deleted 100644 → 0
+ 0
− 160
View file @ f833d2c9
from __future__ import print_function, division
import warnings
import numpy as np
import sympy
from .algorithms import read_coordinates
from .algorithms import discretize
from .postprocessing import offset_to_direction
from .postprocessing import make_kwant_functions
from .postprocessing import offset_to_direction
from .interpolation import interpolate_tb_hamiltonian
try:
# normal situation
from kwant import Builder
from kwant import TranslationalSymmetry
from kwant import HoppingKind
from kwant.lattice import Monatomic
except ImportError:
# probably run on gitlab-ci
pass
class Discretizer(object):
"""Discretize continous Hamiltonian into its tight binding representation.
This class provides easy and nice interface for passing models to Kwant.
Parameters:
-----------
hamiltonian : sympy.Expr or sympy.Matrix instance
Symbolic representation of a continous Hamiltonian. Momentum operators
should be taken from ``discretizer.momentum_operators``.
discrete_coordinates : set of strings
Set of coordinates for which momentum operators will be treated as
differential operators. For example ``discrete_coordinates={'x', 'y'}``.
If left as a None they will be obtained from the input hamiltonian by
reading present coordinates and momentum operators.
interpolate : bool
If True all space dependent parameters in onsite and hopping will be
interpolated to depenend only on the values at site positions.
Default is False.
both_hoppings_directions : bool
If True all hoppings will be returned. For example, if set to True, both
hoppings into (1, 0) and (-1, 0) will be returned. Default is False.
verbose : bool
If True additional information will be printed. Default is False.
Attributes:
-----------
symbolic_hamiltonian : dictionary
Dictionary containing symbolic result of discretization. Key is the
direction of the hopping (zeros for onsite)
lattice : kwant.lattice.Monatomic instance
Lattice to create kwant system. Lattice constant is set to
lattice_constant value.
onsite : function
The value of the onsite Hamiltonian.
hoppings : dict
A dictionary with keys being tuples of the lattice hopping, and values
the corresponding value functions.
discrete_coordinates : set of strings
As in input.
input_hamiltonian : sympy.Expr or sympy.Matrix instance
The input hamiltonian after preprocessing (substitution of functions).
"""
def __init__(self, hamiltonian, discrete_coordinates=None,
lattice_constant=1, interpolate=False,
both_hoppings_directions=False, verbose=False):
self.input_hamiltonian = hamiltonian
if discrete_coordinates is None:
self.discrete_coordinates = read_coordinates(hamiltonian)
else:
self.discrete_coordinates = discrete_coordinates
if verbose:
print('Discrete coordinates set to: ',
sorted(self.discrete_coordinates), end='\n\n')
# discretization
if self.discrete_coordinates:
tb_ham = discretize(hamiltonian, self.discrete_coordinates)
tb_ham = offset_to_direction(tb_ham, self.discrete_coordinates)
else:
tb_ham = {(0,0,0): hamiltonian}
self.discrete_coordinates = {'x', 'y', 'z'}
if interpolate:
tb_ham = interpolate_tb_hamiltonian(tb_ham)
if not both_hoppings_directions:
keys = list(tb_ham)
tb_ham = {k: v for k, v in tb_ham.items()
if k in sorted(keys)[len(keys)//2:]}
self.symbolic_hamiltonian = tb_ham.copy()
for key, val in tb_ham.items():
tb_ham[key] = val.subs(sympy.Symbol('a'), lattice_constant)
# making kwant lattice
dim = len(self.discrete_coordinates)
self.lattice = Monatomic(lattice_constant*np.eye(dim).reshape(dim, dim))
self.lattice_constant = lattice_constant
# making kwant functions
tb = make_kwant_functions(tb_ham, self.discrete_coordinates, verbose)
self.onsite = tb.pop((0,)*len(self.discrete_coordinates))
self.hoppings = {HoppingKind(d, self.lattice): val
for d, val in tb.items()}
def build(self, shape, start, symmetry=None, periods=None):
"""Build Kwant's system.
Convienient functions that simplifies building of a Kwant's system.
Parameters:
-----------
shape : function
A function of real space coordinates that returns a truth value:
true for coordinates inside the shape, and false otherwise.
start : 1d array-like
The real-space origin for the flood-fill algorithm.
symmetry : 1d array-like
Deprecated. Please use ```periods=[symmetry]`` instead.
periods : list of tuples
If periods are provided a translational invariant system will be
built. Periods corresponds basically to a translational symmetry
defined in real space. This vector will be scalled by a lattice
constant before passing it to ``kwant.TranslationalSymmetry``.
Examples: ``periods=[(1,0,0)]`` or ``periods=[(1,0), (0,1)]``.
In second case one will need https://gitlab.kwant-project.org/cwg/wraparound
in order to finalize system.
Returns:
--------
system : kwant.Builder instance
"""
if symmetry is not None:
warnings.warn("\nSymmetry argument is deprecated. " +
"Please use ```periods=[symmetry]`` instead.",
DeprecationWarning)
periods = [symmetry]
if periods is None:
sys = Builder()
else:
vecs = [self.lattice.vec(p) for p in periods]
sys = Builder(TranslationalSymmetry(*vecs))
sys[self.lattice.shape(shape, start)] = self.onsite
for hop, val in self.hoppings.items():
sys[hop] = val
return sys
discretizer/interpolation.py deleted 100644 → 0
+ 0
− 109
View file @ f833d2c9
from __future__ import print_function, division
import itertools
import numpy as np
import sympy
def _follow_path(expr, path):
res = expr
for i in np.arange(len(path)):
res = res.args[path[i]]
return res
def _interchange(expr, sub, path):
res = sub
for i in np.arange(len(path)):
temp = _follow_path(expr, path[:-(i+1)])
args = list(temp.args)
args[path[len(path)-i-1]] = res
res = temp.func(*tuple(args))
return res
def _interpolate_Function(expr):
path = 'None'
factor = 'None'
change = False
summand_0 = 'None'
summand_1 = 'None'
res = expr
for i in np.arange(len(expr.args)):
argument = sympy.expand(expr.args[i])
if argument.func == sympy.Add:
for j in np.arange(len(argument.args)):
summand = argument.args[j]
if summand.func == sympy.Mul:
for k in np.arange(len(summand.args)):
temp = 0
if summand.args[k] == sympy.Symbol('a'):
temp = sympy.Mul(sympy.Mul(*summand.args[:k]),
sympy.Mul(*summand.args[k+1:]))
#print(temp)
if not temp == int(temp):
#print('found one')
factor = (temp)
path = np.array([i, j, k])
if not factor == 'None':
change = True
sign = np.sign(factor)
offsets = np.array([int(factor), sign * (int(sign * factor) + 1)])
weights = 1/np.abs(offsets - factor)
weights = weights/np.sum(weights)
res = ( weights[0] * _interchange(expr, offsets[0] * sympy.Symbol('a'), path[:-1])
+ weights[1] * _interchange(expr, offsets[1] * sympy.Symbol('a'), path[:-1]))
return sympy.expand(res), change
def _interpolate_expression(expr):
change = False
expr = sympy.expand(expr)
res = expr
if isinstance(expr, sympy.Function):# and not change:
path = np.array([])
temp, change = _interpolate_Function(_follow_path(expr, path))
res = _interchange(expr, temp, path)
for i in np.arange(len(expr.args)):
path = np.array([i])
if isinstance(_follow_path(expr, path), sympy.Function) and not change:
temp, change = _interpolate_Function(_follow_path(expr, path))
res = _interchange(expr, temp, path)
for j in np.arange(len(expr.args[i].args)):
path = np.array([i, j])
if isinstance(_follow_path(expr, path), sympy.Function) and not change:
temp, change = _interpolate_Function(_follow_path(expr, path))
res = _interchange(expr, temp, path)
if change:
res = _interpolate_expression(res)
return sympy.expand(res)
def _interpolate(expression):
if not isinstance(expression, sympy.Matrix):
return _interpolate_expression(expression)
shape = expression.shape
interpolated = sympy.zeros(*shape)
for i,j in itertools.product(range(shape[0]), repeat=2):
interpolated[i,j] = _interpolate_expression(expression[i, j])
return interpolated
def interpolate_tb_hamiltonian(tb_hamiltonian):
"""Interpolate tight binding hamiltonian.
This function perform linear interpolation to provide onsite and hoppings
depending only on parameters values at sites positions.
"""
interpolated = {}
for key, val in tb_hamiltonian.items():
interpolated[key] = _interpolate(val)
return interpolated
discretizer/postprocessing.py deleted 100644 → 0
+ 0
− 214
View file @ f833d2c9
from __future__ import print_function, division
import sympy
from sympy.utilities.lambdify import lambdastr
from sympy.printing.lambdarepr import LambdaPrinter
from sympy.core.function import AppliedUndef
class NumericPrinter(LambdaPrinter):
def _print_ImaginaryUnit(self, expr):
return "1.j"
def offset_to_direction(discrete_hamiltonian, discrete_coordinates):
"""Translate hopping keys from offsets to directions.
Parameters:
-----------
discrete_hamiltonian: dict
Discretized hamiltonian, key should be an offset of a hopping and value
corresponds to symbolic hopping.
discrete_coordinates: set
Set of discrete coordinates.
Returns:
--------
discrete_hamiltonian: dict
Discretized hamiltonian, key is a direction of a hopping and value
corresponds to symbolic hopping.
Note:
-----
Coordinates (x,y,z) in output stands for a position of a source of the
hopping.
"""
coordinates = sorted(list(discrete_coordinates))
coordinates = [sympy.Symbol(s, commutative=False) for s in coordinates]
a = sympy.Symbol('a')
onsite_zeros = (0,)*len(discrete_coordinates)
output = {onsite_zeros: discrete_hamiltonian.pop(onsite_zeros)}
for offset, hopping in discrete_hamiltonian.items():
direction = tuple(-c for c in offset)
subs = {c: c + d*a for c, d in zip(coordinates, direction)}
output[direction] = hopping.subs(subs)
return output
# ************ Making kwant functions ***********
def make_return_string(expr):
"""Process a sympy expression into an evaluatable Python return statement.
Parameters:
-----------
expr : sympy.Expr instance
Returns:
--------
output : string
A return string that can be used to assemble a Kwant value function.
func_symbols : set of sympy.Symbol instances
All space dependent functions that appear in the expression.
const_symbols : set of sympy.Symbol instances
All constants that appear in the expression.
"""
func_symbols = {sympy.Symbol(i.func.__name__) for i in
expr.atoms(AppliedUndef)}
free_symbols = {i for i in expr.free_symbols if i.name not in ['x', 'y', 'z']}
const_symbols = free_symbols - func_symbols
expr = expr.subs(sympy.I, sympy.Symbol('1.j')) # quick hack
output = lambdastr((), expr, printer=NumericPrinter)[len('lambda : '):]
output = output.replace('MutableDenseMatrix', 'np.array')
output = output.replace('ImmutableMatrix', 'np.array')
return 'return {}'.format(output), func_symbols, const_symbols
def assign_symbols(func_symbols, const_symbols, discrete_coordinates,
onsite=True):
"""Generate a series of assingments defining a set of symbols.
Parameters:
-----------
func_symbols : set of sympy.Symbol instances
All space dependent functions that appear in the expression.
const_symbols : set of sympy.Symbol instances
All constants that appear in the expression.
Returns:
--------
assignments : list of strings
List of lines used for including in a function.
Notes:
where A, B, C are all the free symbols plus the symbols that appear on the
------
The resulting lines begin with a coordinates assignment of form
`x,y,z = site.pos` when onsite=True, or
`x,y,z = site2.pos` when onsite=False
followed by two lines of form
`A, B, C = p.A, p.B, p.C`
`f, g, h = p.f, p.g, p.h`
where A, B, C are symbols representing constants and f, g, h are symbols
representing functions. Separation of constant and func symbols is probably
not necessary but I leave it for now, just in case.
"""
lines = []
func_names = [i.name for i in func_symbols]
const_names = [i.name for i in const_symbols]
if func_names:
lines.insert(0, ', '.join(func_names) + ' = p.' +
', p.'.join(func_names))
if const_names:
lines.insert(0, ', '.join(const_names) + ' = p.' +
', p.'.join(const_names))
if onsite:
site = 'site'
else:
site = 'site2'
names = sorted(list(discrete_coordinates))
lines.insert(0, '({}, ) = {}.pos'.format(', '.join(names), site))
return lines
def value_function(content, name='_anonymous_func', onsite=True, verbose=False):
"""Generate a Kwant value function from a list of lines containing its body.
Parameters:
-----------
content : list of lines
Lines forming the body of the function.
name : string
Function name (not important).
onsite : bool
If True, the function call signature will be `f(site, p)`, otherwise
`f(site1, site2, p)`.
verbose : bool
Whether the function bodies should be printed.
Returns:
--------
f : function
The function defined in a separated namespace.
"""
if not content[-1].startswith('return'):
raise ValueError('The function does not end with a return statement')
separator = '\n' + 4 * ' '
site_string = 'site' if onsite else 'site1, site2'
header = 'def {0}({1}, p):'.format(name, site_string)
func_code = separator.join([header] + list(content))
namespace = {}
if verbose:
print(func_code)
exec("from __future__ import division", namespace)
exec("import numpy as np", namespace)
exec("from numpy import *", namespace)
exec(func_code, namespace)
return namespace[name]
def make_kwant_functions(discrete_hamiltonian, discrete_coordinates,
verbose=False):
"""Transform discrete hamiltonian into valid kwant functions.
Parameters:
-----------
discrete_hamiltonian: dict
dict in which key is offset of hopping ((0, 0, 0) for onsite)
and value is corresponding symbolic hopping (onsite).
verbose : bool
Whether the function bodies should be printed.
discrete_coordinates : tuple/list
List of discrete coordinates. Must corresponds to offsets in
discrete_hamiltonian keys.
Note:
-----
"""
dim = len(discrete_coordinates)
if not all(len(i)==dim for i in list(discrete_hamiltonian.keys())):
raise ValueError("Dimension of offsets and discrete_coordinates" +
"do not match.")
functions = {}
for offset, hopping in discrete_hamiltonian.items():
onsite = True if all(i == 0 for i in offset) else False
return_string, func_symbols, const_symbols = make_return_string(hopping)
lines = assign_symbols(func_symbols, const_symbols, onsite=onsite,
discrete_coordinates=discrete_coordinates)
lines.append(return_string)
if verbose:
print("Function generated for {}:".format(offset))
f = value_function(lines, verbose=verbose, onsite=onsite)
print()
else:
f = value_function(lines, verbose=verbose, onsite=onsite)
functions[offset] = f
return functions
discretizer/tests/test_algorithms.py deleted 100644 → 0
+ 0
− 120
View file @ f833d2c9
from __future__ import print_function, division
import sympy
import discretizer
from discretizer.algorithms import read_coordinates
from discretizer.algorithms import split_factors
from discretizer.algorithms import wavefunction_name
from discretizer.algorithms import derivate
from discretizer.algorithms import _discretize_summand
from nose.tools import raises
from nose.tools import assert_raises
import numpy as np
kx, ky, kz = sympy.symbols('k_x k_y k_z', commutative=False)
x, y, z = sympy.symbols('x y z', commutative=False)
ax, ay, az = sympy.symbols('a_x a_y a_z')
wf = sympy.Function(wavefunction_name)
Psi = sympy.Function(wavefunction_name)(x, y, z)
A, B = sympy.symbols('A B', commutative=False)
ns = {'A': A, 'B': B, 'a_x': ax, 'a_y': ay, 'az': az, 'x': x, 'y': y, 'z': z}
def test_read_coordinates():
test = {
kx**2 : {'x'},
kx**2 + ky**2 : {'x', 'y'},
kx**2 + ky**2 + kz**2 : {'x', 'y', 'z'},
ky**2 + kz**2 : {'y', 'z'},
kz**2 : {'z'},
kx * A(x,y) * kx : {'x', 'y'},
kx**2 + kz * B(y) : {'x', 'y', 'z'},
}
for inp, out in test.items():
got = read_coordinates(inp)
assert got == out,\
"Should be: split_factors({})=={}. Not {}".format(inp, out, got)
def test_split_factors_1():
test = {
kz * Psi : (1, kz, Psi),
A * kx**2 * Psi : (A * kx, kx, Psi),
A * kx**2 * ky * Psi : (A * kx**2, ky, Psi),
ky * A * kx * B * Psi : (ky * A, kx, B * Psi),
kx : (1, kx, 1),
kx**2 : (kx, kx, 1),
A : (1, 1, A),
A**2 : (1, 1, A**2),
kx*A**2 : (1, kx, A**2),
kx**2*A**2 : (kx, kx, A**2),
A(x, y, z) : (1, 1, A(x, y, z)),
Psi : (1, 1, Psi),
np.int(5) : (1, 1, np.int(5)),
np.float(5) : (1, 1, np.float(5)),
sympy.Integer(5) : (1, 1, sympy.Integer(5)),
sympy.Float(5) : (1, 1, sympy.Float(5)),
1 : (1, 1, 1),
1.0 : (1, 1, 1.0),
5 : (1, 1, 5),
5.0 : (1, 1, 5.0),
}
for inp, out in test.items():
got = split_factors(inp, discrete_coordinates={'x', 'y', 'z'})
assert got == out,\
"Should be: split_factors({})=={}. Not {}".format(inp, out, got)
@raises(AssertionError)
def test_split_factors_2():
split_factors(A+B, discrete_coordinates={'x', 'y', 'z'})
def test_derivate_1():
test = {
(A(x), kx): '-I*(-A(-a_x + x)/(2*a_x) + A(a_x + x)/(2*a_x))',
(A(x), ky): '0',
(A(x)*B, kx): '-I*(-A(-a_x + x)*B/(2*a_x) + A(a_x + x)*B/(2*a_x))',
(A(x) + B(x), kx): '-I*(-A(-a_x + x)/(2*a_x) + A(a_x + x)/(2*a_x) - B(-a_x + x)/(2*a_x) + B(a_x + x)/(2*a_x))',
(A, kx): '0',
(5, kx): '0',
(A(x) * B(x), kx): '-I*(-A(-a_x + x)*B(-a_x + x)/(2*a_x) + A(a_x + x)*B(a_x + x)/(2*a_x))',
(Psi, ky): '-I*(-Psi(x, -a_y + y, z)/(2*a_y) + Psi(x, a_y + y, z)/(2*a_y))',
}
for inp, out in test.items():
got = (derivate(*inp))
out = sympy.sympify(out, locals=ns)
assert sympy.simplify(sympy.expand(got - out)) == 0,\
"Should be: derivate({0[0]}, {0[1]})=={1}. Not {2}".format(inp, out, got)
@raises(TypeError)
def test_derivate_2():
derivate(A(x), kx**2)
@raises(ValueError)
def test_derivate_3():
derivate(A(x), sympy.Symbol('A'))
def test_discretize_summand_1():
test = {
kx * A(x): '-I*(-A(-a_x + x)/(2*a_x) + A(a_x + x)/(2*a_x))',
kx * Psi: '-I*(-Psi(-a_x + x, y, z)/(2*a_x) + Psi(a_x + x, y, z)/(2*a_x))',
kx**2 * Psi: 'Psi(x, y, z)/(2*a_x**2) - Psi(-2*a_x + x, y, z)/(4*a_x**2) - Psi(2*a_x + x, y, z)/(4*a_x**2)',
kx * A(x) * kx * Psi: 'A(-a_x + x)*Psi(x, y, z)/(4*a_x**2) - A(-a_x + x)*Psi(-2*a_x + x, y, z)/(4*a_x**2) + A(a_x + x)*Psi(x, y, z)/(4*a_x**2) - A(a_x + x)*Psi(2*a_x + x, y, z)/(4*a_x**2)',
}
for inp, out in test.items():
got = _discretize_summand(inp, discrete_coordinates={'x', 'y', 'z'})
out = sympy.sympify(out, locals=ns)
assert sympy.simplify(sympy.expand(got - out)) == 0,\
"Should be: _discretize_summand({})=={}. Not {}".format(inp, out, got)
@raises(AssertionError)
def test_discretize_summand_2():
_discretize_summand(kx*A(x)+ B(x), discrete_coordinates={'x', 'y', 'z'})
discretizer/tests/test_hoppings.py deleted 100644 → 0
+ 0
− 129
View file @ f833d2c9
from __future__ import print_function, division
import sympy
import discretizer
from discretizer.algorithms import wavefunction_name
from discretizer.algorithms import read_hopping_from_wf
from discretizer.algorithms import extract_hoppings
from nose.tools import raises
from nose.tools import assert_raises
import numpy as np
kx, ky, kz = sympy.symbols('k_x k_y k_z', commutative=False)
x, y, z = sympy.symbols('x y z', commutative=False)
ax, ay, az = sympy.symbols('a_x a_y a_z')
wf = sympy.Function(wavefunction_name)
Psi = sympy.Function(wavefunction_name)(x, y, z)
A, B = sympy.symbols('A B', commutative=False)
ns = {'A': A, 'B': B, 'a_x': ax, 'a_y': ay, 'az': az, 'x': x, 'y': y, 'z': z}
def test_read_hoppings_from_wf_1():
offsets = [(0,0,0), (1,0,0), (0,1,0), (0,0,1), (1,1,1), (1,2,3)]
test = {}
for offset in offsets:
nx, ny, nz = offset
key = Psi.subs({x: x + nx*ax, y: y + ny * ay, z: z + nz * az})
test[key] = offset
for inp, out in test.items():
got = read_hopping_from_wf(inp)
assert got == out,\
"Should be: read_hopping_from_wf({}) == {}. Not {}".format(inp, out, got)
def test_read_hoppings_from_wf_2():
test = {
wf(x, y, z): (0,0,0),
wf(x, y): (0, 0),
wf(x, z): (0, 0),
wf(x+ax, y-2*ay): (1, -2),
wf(x, z+3*az): (0, 3),
wf(y, z): (0, 0),
wf(y, z+az): (0, 1),
}
for inp, out in test.items():
got = read_hopping_from_wf(inp)
assert got == out,\
"Should be: read_hopping_from_wf({}) == {}. Not {}".format(inp, out, got)
def test_test_read_hoppings_from_wf_ValueError():
tests = {
wf(x+ay, ay),
wf(y, x),
wf(x, x),
wf(x, ax),
wf(y, A),
}
for inp in tests:
assert_raises(ValueError, read_hopping_from_wf, inp)
def test_test_read_hoppings_from_wf_TypeError():
tests = {
wf(x,y,z) + A,
A(x,y),
5*Psi,
Psi+2,
A,
}
for inp in tests:
assert_raises(TypeError, read_hopping_from_wf, inp)
def test_extract_hoppings():
discrete_coordinates = {'x', 'y', 'z'}
tests = [
{
'test_inp': '-I*(-Psi(-a_x + x, y)/(2*a_x) + Psi(a_x + x, y)/(2*a_x))',
'test_out': {
(1, 0): '-I/(2*a)',
(-1, 0): 'I/(2*a)',
},
},
{
'test_inp': 'Psi(x, y)/(2*a_x**2) - Psi(-2*a_x + x, y)/(4*a_x**2) - Psi(2*a_x + x, y)/(4*a_x**2)',
'test_out': {
(1, 0): '-1/a**2',
(0, 0): '2/a**2',
(-1, 0): '-1/a**2',
},
},
{
'test_inp': 'A(-a_x + x)*Psi(x, y)/(4*a_x**2) - A(-a_x + x)*Psi(-2*a_x + x, y)/(4*a_x**2) + A(a_x + x)*Psi(x, y)/(4*a_x**2) - A(a_x + x)*Psi(2*a_x + x, y)/(4*a_x**2)',
'test_out': {
(1, 0): '-A(a/2 + x)/a**2',
(0, 0): 'A(-a/2 + x)/a**2 + A(a/2 + x)/a**2',
(-1, 0): '-A(-a/2 + x)/a**2',
},
},
{
'test_inp': 'I*A(x, y)*Psi(x, -a_y + y, z)/(4*a_x**2*a_y) - I*A(x, y)*Psi(x, a_y + y, z)/(4*a_x**2*a_y) - I*A(-2*a_x + x, y)*Psi(-2*a_x + x, -a_y + y, z)/(8*a_x**2*a_y) + I*A(-2*a_x + x, y)*Psi(-2*a_x + x, a_y + y, z)/(8*a_x**2*a_y) - I*A(2*a_x + x, y)*Psi(2*a_x + x, -a_y + y, z)/(8*a_x**2*a_y) + I*A(2*a_x + x, y)*Psi(2*a_x + x, a_y + y, z)/(8*a_x**2*a_y)',
'test_out': {
(-1, 1, 0): 'I*A(-a + x, y)/(2*a**3)',
(1, 1, 0): 'I*A(a + x, y)/(2*a**3)',
(0, -1, 0): 'I*A(x, y)/a**3',
(0, 1, 0): '-I*A(x, y)/a**3',
(1, -1, 0): '-I*A(a + x, y)/(2*a**3)',
(-1, -1, 0): '-I*A(-a + x, y)/(2*a**3)',
},
},
]
for test in tests:
test_inp = test['test_inp']
test_out = test['test_out']
inp = sympy.sympify(test_inp, locals=ns)
result = extract_hoppings(inp, discrete_coordinates)
for key, got in result.items():
out = sympy.sympify(test_out[key], locals=ns)
assert sympy.simplify(sympy.expand(got - out)) == 0, \
"Should be: extract_hoppings({})=={}. Not {}".format(inp, test_out, result)
kwant/__init__.py
+ 10
− 0
View file @ 6009151a
...@@ -9,6 +9,7 @@ ...@@ -9,6 +9,7 @@
__all__ = [] __all__ = []
import numpy # Needed by C. Gohlke's Windows package. import numpy # Needed by C. Gohlke's Windows package.
import warnings
try: try:
from . import _system from . import _system
...@@ -62,3 +63,12 @@ def test(verbose=True): ...@@ -62,3 +63,12 @@ def test(verbose=True):
"-s"] + (['-v'] if verbose else [])) "-s"] + (['-v'] if verbose else []))
test.__test__ = False test.__test__ = False
# Exposing discretizer
try:
from .continuum import discretize
__all__.extend(['discretize'])
except ImportError:
warnings.warn('Discretizer module not available. Is sympy installed?',
RuntimeWarning)
kwant/continuum/__init__.py 0 → 100644
+ 7
− 0
View file @ 6009151a
__all__ = ['discretize', 'discretize_symbolic', 'build_discretized',
'momentum_operators', 'position_operators']
from .discretizer import discretize, discretize_symbolic, build_discretized
from ._common import momentum_operators, position_operators
from ._common import sympify, lambdify, make_commutative
kwant/continuum/_common.py 0 → 100644
+ 217
− 0
View file @ 6009151a
import numpy as np
import sympy
from sympy.core.function import AppliedUndef
from sympy.core.sympify import converter
from sympy.abc import _clash
import functools
import inspect
from collections import defaultdict
from operator import mul
from sympy.physics.matrices import msigma as _msigma
from sympy.physics.quantum import TensorProduct
momentum_operators = sympy.symbols('k_x k_y k_z', commutative=False)
position_operators = sympy.symbols('x y z', commutative=False)
msigma = lambda i: sympy.eye(2) if i == 0 else _msigma(i)
pauli = (msigma(0), msigma(1), msigma(2), msigma(3))
_clash = _clash.copy()
_clash.update({s.name: s for s in momentum_operators})
_clash.update({s.name: s for s in position_operators})
_clash.update({'kron': TensorProduct, 'eye': sympy.eye})
_clash.update({'sigma_{}'.format(c): p for c, p in zip('0xyz', pauli)})
# workaroud for https://github.com/sympy/sympy/issues/12060
del _clash['I']
del _clash['pi']
################ Various helpers to handle sympy
def lambdify(hamiltonian, *, substitutions=None):
"""Return a callable object for computing continuum Hamiltonian.
Parameters
----------
hamiltonian : sympy.Expr or sympy.Matrix, or string
Symbolic representation of a continous Hamiltonian. When providing
a sympy expression, it is recommended to use operators
defined in `~kwant.continuum.momentum_operators` in order to provide
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]]})``.
"""
expr = hamiltonian
if not isinstance(expr, (sympy.Expr, sympy.matrices.MatrixBase)):
expr = sympify(expr, substitutions)
args = [s.name for s in expr.atoms(sympy.Symbol)]
args += [str(f.func) for f in expr.atoms(AppliedUndef, sympy.Function)]
f = sympy.lambdify(sorted(args), expr)
sig = inspect.signature(f)
pars = list(sig.parameters.values())
pars = [p.replace(kind=inspect.Parameter.KEYWORD_ONLY) for p in pars]
f.__signature__ = inspect.Signature(pars)
return f
def sympify(string, substitutions=None, **kwargs):
"""Return sympified object with respect to kwant-specific rules.
This is a modification of ``sympy.sympify`` to apply kwant-specific rules,
which includes preservation of proper commutation relations between
position and momentum operators, and array to matrix casting.
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.
substitutions : dict, defaults to empty
A namespace to be passed internally to ``sympy.sympify``.
Could be used to 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``.
Example
-------
>>> from kwant.continuum import sympify
>>> hamiltonian = sympify('k_x * A(x) * k_x + V(x)')
"""
stored_value = None
if substitutions is None:
substitutions = {}
substitutions.update(_clash)
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)
finally:
if stored_value is not None:
converter[list] = stored_value
else:
del converter[list]
return sympy.expand(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]
subs = {s: sympy.Symbol(s.name) for s in symbols}
expr = expr.subs(subs)
return expr
def expression_monomials(expression, *gens):
"""Parse ``expression`` into monomials in the symbols in ``gens``.
Example
-------
>>> expr = A * (x**2 + y) + B * x + C
>>> expression_monomials(expr, x, y)
{1: C, x**2: A, y: A, x: B}
"""
if expression.atoms(AppliedUndef):
raise NotImplementedError('Getting monomials of expressions containing '
'functions is not implemented.')
expression = make_commutative(expression, *gens)
gens = [make_commutative(g, g) for g in gens]
expression = sympy.expand(expression)
summands = expression.as_ordered_terms()
output = defaultdict(int)
for summand in summands:
key = [sympy.Integer(1)]
if summand in gens:
key.append(summand)
elif isinstance(summand, sympy.Pow):
if summand.args[0] in gens:
key.append(summand)
else:
for arg in summand.args:
if arg in gens:
key.append(arg)
if isinstance(arg, sympy.Pow):
if arg.args[0] in gens:
key.append(arg)
key = functools.reduce(mul, key)
val = summand.xreplace({g: 1 for g in gens})
### to not create key
if val != 0:
output[key] += val
new_expression = sum(k * v for k, v in output.items())
assert sympy.expand(expression) == sympy.expand(new_expression)
return dict(output)
def matrix_monomials(matrix, *gens):
output = defaultdict(lambda: sympy.zeros(*matrix.shape))
for (i, j), expression in np.ndenumerate(matrix):
summands = expression_monomials(expression, *gens)
for key, val in summands.items():
output[key][i, j] += 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)
kwant/continuum/discretizer.py 0 → 100644
+ 555
− 0
View file @ 6009151a
This diff is collapsed.
kwant/continuum/tests/__init__.py 0 → 100644
+ 0
− 0
View file @ 6009151a
kwant/continuum/tests/test_common.py 0 → 100644
+ 146
− 0
View file @ 6009151a
from kwant.continuum._common import position_operators, momentum_operators
from kwant.continuum._common import make_commutative, sympify
from kwant.continuum._common import expression_monomials, matrix_monomials
from kwant.continuum._common import lambdify
import tinyarray as ta
from sympy.physics.matrices import msigma
import sympy
from functools import reduce
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
# 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']
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]])
A, B, non_x = sympy.symbols('A B x', commutative=False)
x = sympy.Symbol('x')
expr1 = non_x*A*non_x + x**2 * A * x + B*non_x**2
matr = sympy.Matrix([[expr1, expr1+A*non_x], [0, -expr1]])
res_mat = sympy.Matrix([[x**3*A + x**2*A + x**2*B, x**3*A + x**2*A + x**2*B + x*A],
[0, -x**3*A - x**2*A - x**2*B]])
def test_make_commutative():
assert make_commutative(expr1, x) == make_commutative(expr1, non_x)
assert make_commutative(expr1, x) == x**3*A + x**2*A + x**2*B
assert make_commutative(matr, x) == res_mat
expr2 = non_x*A*non_x + x**2 * A*2 * x + B*non_x/2 + non_x*B/2 + x + A + non_x + x/A
def test_expression_monomials():
assert expression_monomials(expr2, x) == {x**3: 2*A, 1: A, x: 2 + A**(-1) + B, x**2: A}
assert expression_monomials(expr1, x) == {x**2: A + B, x**3: A}
assert expression_monomials(x, x) == {x: 1}
assert expression_monomials(x**2, x) == {x**2: 1}
assert expression_monomials(x**2 + x, x) == {x: 1, x**2: 1}
expr = 1 + x + A*x + 2*x + x**2 + A*x**2 + non_x*A*non_x
out = {1: 1, x: 3 + A, x**2: 2 * A + 1}
assert expression_monomials(expr, x) == out
expr = 1 + x * (3 + A) + x**2 * (1 + A)
out = {1: 1, x: 3 + A, x**2: 1 * A + 1}
assert expression_monomials(expr, x) == out
def legacy_expression_monomials(expr, *gens):
""" This was my first implementation. Unfortunately it is very slow.
It is used to test correctness of new matrix_monomials function.
"""
expr = make_commutative(expr, x)
R = sympy.ring(gens, sympy.EX, sympy.lex)[0]
expr = R(expr)
output = {}
for power, coeff in zip(expr.monoms(), expr.coeffs()):
key = reduce(mul, [sympy.Symbol(k.name)**n for k, n in zip(gens, power)])
output[key] = sympy.expand(coeff.as_expr())
return output
def test_expression_monomials_with_reference_function():
assert legacy_expression_monomials(expr2, x) == expression_monomials(expr2, x)
def test_matrix_monomials():
out = {
x**2: sympy.Matrix([[A + B, A + B],[0, -A - B]]),
x: sympy.Matrix([[0, A], [0, 0]]),
x**3: sympy.Matrix([[A, A], [0, -A]])}
mons = matrix_monomials(matr, x)
assert mons == out
@pytest.mark.parametrize("e, should_be, kwargs", [
("x+y", lambda x, y: x+y, dict(x=1, y=2)),
("1", lambda: 1, dict()),
("f(x)", lambda f, x: f(x), dict(f=lambda x: x, x=2)),
(sympify("f(x)"), lambda f, x: f(x), dict(f=lambda x: x, x=2)),
("[[f(x)]]", lambda f, x: ta.array(f(x)), dict(f=lambda x: x, x=2))
])
def test_lambdify(e, should_be, kwargs):
e = lambdify(e)
assert e(**kwargs) == should_be(**kwargs)
# dispersion_string = ('A * k_x**2 * eye(2) + B * k_y**2 * eye(2)'
# '+ alpha * k_x * sigma_y')
# def dispersion_function(kx, ky, A, B, alpha, **kwargs):
# h_k = ((A * kx**2 + B * ky**2) * sigma_0 +
# alpha * kx * sigma_y)
# return np.linalg.eigvalsh(h_k)
# dispersion_kwargs = {'A': 1, 'B': 2, 'alpha': 0.5, 'M': 2}
# sigma_0 = np.eye(2)
# sigma_y = np.array([[0, -1j], [1j, 0]])
# @pytest.mark.parametrize("expr, should_be, kwargs", [
# (dispersion_string, dispersion_function, dispersion_kwargs),
# (sympify(dispersion_string), dispersion_function, dispersion_kwargs),
# ])
# def test_lambdify(expr, should_be, kwargs):
# N = 5
# x = y = np.linspace(-N, N+1)
# xy = list(itertools.product(x, y))
# h_k = lambdify(expr)
# energies = [la.eigvalsh(h_k(k_x=kx, k_y=ky, **kwargs))
# for kx, ky in xy]
# energies_should_be = [should_be(kx, ky, **kwargs) for kx, ky in xy]
# assert np.allclose(energies, energies_should_be)
kwant/continuum/tests/test_discretizer.py 0 → 100644
+ 480
− 0
View file @ 6009151a
import sympy
import numpy as np
from ..discretizer import discretize
from ..discretizer import discretize_symbolic
from ..discretizer import build_discretized
from ..discretizer import _wf
import inspect
from functools import wraps
def swallows_extra_kwargs(f):
sig = inspect.signature(f)
pars = sig.parameters
if any(i.kind is inspect.Parameter.VAR_KEYWORD for i in pars.values()):
return f
names = {name for name, value in pars.items() if
value.kind not in (inspect.Parameter.VAR_POSITIONAL,
inspect.Parameter.POSITIONAL_ONLY)}
@wraps(f)
def wrapped(*args, **kwargs):
bound = sig.bind(*args, **{name: value for name, value
in kwargs.items() if name in names})
return f(*bound.args, **bound.kwargs)
return wrapped
I = sympy.I
kx, ky, kz = sympy.symbols('k_x k_y k_z', commutative=False)
x, y, z = sympy.symbols('x y z', commutative=False)
ax, ay, az = sympy.symbols('a_x a_y a_z')
a = sympy.symbols('a')
wf = _wf
Psi = wf(x, y, z)
A, B = sympy.symbols('A B', commutative=False)
ns = {'A': A, 'B': B, 'a_x': ax, 'a_y': ay, 'az': az, 'x': x, 'y': y, 'z': z}
def test_reading_coordinates():
test = {
kx**2 : {'x'},
kx**2 + ky**2 : {'x', 'y'},
kx**2 + ky**2 + kz**2 : {'x', 'y', 'z'},
ky**2 + kz**2 : {'y', 'z'},
kz**2 : {'z'},
kx * A(x,y) * kx : {'x'},
kx**2 + kz * B(y) : {'x', 'z'},
}
for inp, out in test.items():
ham, got = discretize_symbolic(inp)
assert all(d in out for d in got),\
"Should be: _split_factors({})=={}. Not {}".format(inp, out, got)
def test_reading_coordinates_matrix():
test = [
(sympy.Matrix([kx**2]) , {'x'}),
(sympy.Matrix([kx**2 + ky**2]) , {'x', 'y'}),
(sympy.Matrix([kx**2 + ky**2 + kz**2]) , {'x', 'y', 'z'}),
(sympy.Matrix([ky**2 + kz**2]) , {'y', 'z'}),
(sympy.Matrix([kz**2]) , {'z'}),
(sympy.Matrix([kx * A(x,y) * kx]) , {'x'}),
(sympy.Matrix([kx**2 + kz * B(y)]) , {'x', 'z'}),
]
for inp, out in test:
ham, got = discretize_symbolic(inp)
assert all(d in out for d in got),\
"Should be: _split_factors({})=={}. Not {}".format(inp, out, got)
def test_reading_different_matrix_types():
test = [
(sympy.MutableMatrix([kx**2]) , {'x'}),
(sympy.ImmutableMatrix([kx**2]) , {'x'}),
(sympy.MutableDenseMatrix([kx**2]) , {'x'}),
(sympy.ImmutableDenseMatrix([kx**2]) , {'x'}),
]
for inp, out in test:
ham, got = discretize_symbolic(inp)
assert all(d in out for d in got),\
"Should be: _split_factors({})=={}. Not {}".format(inp, out, got)
def test_simple_derivations():
test = {
kx**2 : {(0,): 2/a**2, (1,): -1/a**2},
kx**2 + ky**2 : {(0, 1): -1/a**2, (0, 0): 4/a**2,
(1, 0): -1/a**2},
kx**2 + ky**2 + kz**2 : {(1, 0, 0): -1/a**2, (0, 0, 1): -1/a**2,
(0, 0, 0): 6/a**2, (0, 1, 0): -1/a**2},
ky**2 + kz**2 : {(0, 1): -1/a**2, (0, 0): 4/a**2,
(1, 0): -1/a**2},
kz**2 : {(0,): 2/a**2, (1,): -1/a**2},
kx * A(x,y) * kx : {(1, ): -A(a/2 + x, y)/a**2,
(0, ): A(-a/2 + x, y)/a**2 + A(a/2 + x, y)/a**2},
kx**2 + kz * B(y) : {(1, 0): -1/a**2, (0, 1): -I*B(y)/(2*a),
(0, 0): 2/a**2},
kx * A(x) : {(0,): 0, (1,): -I*A(a + x)/(2*a)},
ky * A(x) : {(1,): -I*A(x)/(2*a), (0,): 0},
kx * A(x) * B : {(0,): 0, (1,): -I*A(a + x)*B/(2*a)},
5 * kx : {(0,): 0, (1,): -5*I/(2*a)},
kx * (A(x) + B(x)) : {(0,): 0,
(1,): -I*A(a + x)/(2*a) - I*B(a + x)/(2*a)},
}
for inp, out in test.items():
got, _ = discretize_symbolic(inp)
assert got == out
for inp, out in test.items():
got, _ = discretize_symbolic(str(inp), substitutions=ns)
assert got == out
def test_simple_derivations_matrix():
test = {
kx**2 : {(0,): 2/a**2, (1,): -1/a**2},
kx**2 + ky**2 : {(0, 1): -1/a**2, (0, 0): 4/a**2,
(1, 0): -1/a**2},
kx**2 + ky**2 + kz**2 : {(1, 0, 0): -1/a**2, (0, 0, 1): -1/a**2,
(0, 0, 0): 6/a**2, (0, 1, 0): -1/a**2},
ky**2 + kz**2 : {(0, 1): -1/a**2, (0, 0): 4/a**2,
(1, 0): -1/a**2},
kz**2 : {(0,): 2/a**2, (1,): -1/a**2},
kx * A(x,y) * kx : {(1, ): -A(a/2 + x, y)/a**2,
(0, ): A(-a/2 + x, y)/a**2 + A(a/2 + x, y)/a**2},
kx**2 + kz * B(y) : {(1, 0): -1/a**2, (0, 1): -I*B(y)/(2*a),
(0, 0): 2/a**2},
kx * A(x) : {(0,): 0, (1,): -I*A(a + x)/(2*a)},
ky * A(x) : {(1,): -I*A(x)/(2*a), (0,): 0},
kx * A(x) * B : {(0,): 0, (1,): -I*A(a + x)*B/(2*a)},
5 * kx : {(0,): 0, (1,): -5*I/(2*a)},
kx * (A(x) + B(x)) : {(0,): 0,
(1,): -I*A(a + x)/(2*a) - I*B(a + x)/(2*a)},
}
new_test = []
for inp, out in test.items():
new_out = {}
for k, v in out.items():
new_out[k] = sympy.Matrix([v])
new_test.append((sympy.Matrix([inp]), new_out))
for inp, out in new_test:
got, _ = discretize_symbolic(inp)
assert got == out
for inp, out in new_test:
got, _ = discretize_symbolic(str(inp), substitutions=ns)
assert got == out
for inp, out in new_test:
got, _ = discretize_symbolic(str(inp).replace('Matrix', ''), substitutions=ns)
assert got == out
def test_integer_float_input():
test = {
0 : {(0,0,0): 0},
1 : {(0,0,0): 1},
5 : {(0,0,0): 5},
}
for inp, out in test.items():
got, _ = discretize_symbolic(int(inp), {'x', 'y', 'z'})
assert got == out
got, _ = discretize_symbolic(float(inp), {'x', 'y', 'z'})
assert got == out
# let's test in matrix version too
new_test = []
for inp, out in test.items():
new_out = {}
for k, v in out.items():
new_out[k] = sympy.Matrix([v])
new_test.append((inp, new_out))
for inp, out in new_test:
got, _ = discretize_symbolic(sympy.Matrix([int(inp)]), {'x', 'y', 'z'})
assert got == out
got, _ = discretize_symbolic(sympy.Matrix([float(inp)]), {'x', 'y', 'z'})
assert got == out
def test_different_discrete_coordinates():
test = [
(
{'x', 'y', 'z'}, {
(1, 0, 0): -1/a**2, (0, 0, 1): -1/a**2,
(0, 0, 0): 6/a**2, (0, 1, 0): -1/a**2
}
),
(
{'x', 'y'}, {
(0, 1): -1/a**2,
(1, 0): -1/a**2,
(0, 0): kz**2 + 4/a**2
}
),
(
{'x', 'z'}, {
(0, 1): -1/a**2,
(1, 0): -1/a**2,
(0, 0): ky**2 + 4/a**2
}
),
(
{'y', 'z'}, {
(0, 1): -1/a**2,
(1, 0): -1/a**2,
(0, 0): kx**2 + 4/a**2
}
),
(
{'x'}, {
(0,): ky**2 + kz**2 + 2/a**2, (1,): -1/a**2
}
),
(
{'y'}, {
(0,): kx**2 + kz**2 + 2/a**2, (1,): -1/a**2
}
),
(
{'z'}, {
(0,): ky**2 + kx**2 + 2/a**2, (1,): -1/a**2
}
) ,
]
for inp, out in test:
got, _ = discretize_symbolic(kx**2 + ky**2 + kz**2, inp)
assert got == out
# let's test in matrix version too
new_test = []
for inp, out in test:
new_out = {}
for k, v in out.items():
new_out[k] = sympy.Matrix([v])
new_test.append((inp, new_out))
for inp, out in new_test:
got, _ = discretize_symbolic(sympy.Matrix([kx**2 + ky**2 + kz**2]), inp)
assert got == out
def test_non_expended_input():
symbolic, coords = discretize_symbolic(kx * (kx + A(x)))
desired = {
(0,): 2/a**2,
(1,): -I*A(a + x)/(2*a) - 1/a**2
}
assert symbolic == desired
def test_matrix_with_zeros():
Matrix = sympy.Matrix
symbolic, _ = discretize_symbolic("[[k_x*A(x)*k_x, 0], [0, k_x*A(x)*k_x]]")
output = {
(0,) : Matrix([[A(-a/2 + x)/a**2 + A(a/2 + x)/a**2, 0], [0, A(-a/2 + x)/a**2 + A(a/2 + x)/a**2]]),
(1,) : Matrix([[-A(a/2 + x)/a**2, 0], [0, -A(a/2 + x)/a**2]]),
}
assert symbolic == output
def test_numeric_functions_basic_symbolic():
for i in [0, 1, 3, 5]:
builder = discretize(i, {'x'})
lat = next(iter(builder.sites()))[0]
assert builder[lat(0)] == i
p = dict(t=i)
tb = {(0,): sympy.sympify("2*t"), (1,): sympy.sympify('-t')}
builder = build_discretized(tb, {'x'}, lattice_constant=1)
lat = next(iter(builder.sites()))[0]
assert 2*p['t'] == builder[lat(0)](None, **p)
assert -p['t'] == builder[lat(1), lat(0)](None, None, **p)
tb = {(0,): sympy.sympify("0"), (1,): sympy.sympify('-1j * t')}
builder = build_discretized(tb, {'x'}, lattice_constant=1)
lat = next(iter(builder.sites()))[0]
assert -1j * p['t'] == builder[lat(0), lat(1)](None, None, **p)
assert +1j * p['t'] == builder[lat(1), lat(0)](None, None, **p)
def test_numeric_functions_not_discrete_coords():
builder = discretize('k_y + y', 'x')
lat = next(iter(builder.sites()))[0]
onsite = builder[lat(0)]
assert onsite(None, k_y=2, y=1) == 2 + 1
def test_numeric_functions_with_pi():
# Two cases because once it is casted
# to complex, one there is a function created
builder = discretize('A + pi', 'x')
lat = next(iter(builder.sites()))[0]
onsite = builder[lat(0)]
assert onsite(None, A=1) == 1 + np.pi
builder = discretize('pi', 'x')
lat = next(iter(builder.sites()))[0]
onsite = builder[lat(0)]
assert onsite == np.pi
def test_numeric_functions_basic_string():
for i in [0, 1, 3, 5]:
builder = discretize(i, {'x'})
lat = next(iter(builder.sites()))[0]
assert builder[lat(0)] == i
p = dict(t=i)
tb = {(0,): "2*t", (1,): "-t"}
builder = build_discretized(tb, {'x'}, lattice_constant=1)
lat = next(iter(builder.sites()))[0]
assert 2*p['t'] == builder[lat(0)](None, **p)
assert -p['t'] == builder[lat(1), lat(0)](None, None, **p)
tb = {(0,): "0", (1,): "-1j * t"}
builder = build_discretized(tb, {'x'}, lattice_constant=1)
lat = next(iter(builder.sites()))[0]
assert -1j * p['t'] == builder[lat(0), lat(1)](None, None, **p)
assert +1j * p['t'] == builder[lat(1), lat(0)](None, None, **p)
tb = {(0,): "0", (-1,): "+1j * t"}
builder = build_discretized(tb, {'x'}, lattice_constant=1)
lat = next(iter(builder.sites()))[0]
assert -1j * p['t'] == builder[lat(0), lat(1)](None, None, **p)
assert +1j * p['t'] == builder[lat(1), lat(0)](None, None, **p)
def test_numeric_functions_advance():
hams = [
kx**2,
kx**2 + x,
A(x),
kx*A(x)*kx,
sympy.Matrix([[kx * A(x) * kx, A(x)*kx], [kx*A(x), A(x)+B]]),
kx**2 + B * x,
'k_x**2 + sin(x)',
B ** 0.5 * kx**2,
B ** (1/2) * kx**2,
sympy.sqrt(B) * kx**2,
]
for hamiltonian in hams:
for a in [1, 2, 5]:
for fA in [lambda x: x, lambda x: x**2, lambda x: x**3]:
symbolic, coords = discretize_symbolic(hamiltonian, {'x'})
builder = build_discretized(symbolic, coords, lattice_constant=a)
lat = next(iter(builder.sites()))[0]
p = dict(A=fA, B=5, sin=np.sin)
# test onsite
v = symbolic.pop((0,)).subs({sympy.symbols('a'): a, B: p['B']})
f_sym = sympy.lambdify(['A', 'x'], v)
f_num = builder[lat(0)]
if callable(f_num):
f_num = swallows_extra_kwargs(f_num)
for n in range(-100, 100, 10):
assert np.allclose(f_sym(fA, a*n), f_num(lat(n), **p))
else:
for n in range(-100, 100, 10):
assert np.allclose(f_sym(fA, a*n), f_num)
# test hoppings
for k, v in symbolic.items():
v = v.subs({sympy.symbols('a'): a, B: p['B']})
f_sym = sympy.lambdify(['A', 'x'], v)
f_num = builder[lat(0), lat(k[0])]
if callable(f_num):
f_num = swallows_extra_kwargs(f_num)
for n in range(10):
lhs = f_sym(fA, a * n)
rhs = f_num(lat(n), lat(n+k[0]), **p)
assert np.allclose(lhs, rhs)
else:
for n in range(10):
lhs = f_sym(fA, a * n)
rhs = f_num
assert np.allclose(lhs, rhs)
def test_numeric_functions_with_parameter():
hams = [
kx**2 + A(B, x)
]
for hamiltonian in hams:
for a in [1, 2, 5]:
for fA in [lambda c, x: x+c, lambda c, x: x**2 + c]:
symbolic, coords = discretize_symbolic(hamiltonian, {'x'})
builder = build_discretized(symbolic, coords, lattice_constant=a)
lat = next(iter(builder.sites()))[0]
p = dict(A=fA, B=5)
# test onsite
v = symbolic.pop((0,)).subs({sympy.symbols('a'): a, B: p['B']})
f_sym = sympy.lambdify(['A', 'x'], v)
f_num = builder[lat(0)]
if callable(f_num):
f_num = swallows_extra_kwargs(f_num)
for n in range(10):
s = lat(n)
xi = a * n
if callable(f_num):
assert np.allclose(f_sym(fA, xi), f_num(s, **p))
else:
assert np.allclose(f_sym(fA, xi), f_num)
# test hoppings
for k, v in symbolic.items():
v = v.subs({sympy.symbols('a'): a, B: p['B']})
f_sym = sympy.lambdify(['A', 'x'], v)
f_num = builder[lat(0), lat(k[0])]
if callable(f_num):
f_num = swallows_extra_kwargs(f_num)
for n in range(10):
s = lat(n)
xi = a * n
lhs = f_sym(fA, xi)
if callable(f_num):
rhs = f_num(lat(n), lat(n+k[0]), **p)
else:
rhs = f_num
assert np.allclose(lhs, rhs)
def test_basic_verbose(capsys): # or use "capfd" for fd-level
discretize('k_x * A(x) * k_x', verbose=True)
out, err = capsys.readouterr()
assert "Discrete coordinates set to" in out
assert "Function generated for (0,)" in out
def test_that_verbose_covers_all_hoppings(capsys):
discretize('k_x**2 + k_y**2 + k_x*k_y', verbose=True)
out, err = capsys.readouterr()
for tag in [(0, 1), (0, 0), (1, -1), (1, 1)]:
assert "Function generated for {}".format(tag) in out
def test_verbose_cache(capsys):
discretize('[[k_x * A(x) * k_x]]', verbose=True)
out, err = capsys.readouterr()
assert '_cache_0' in out
def test_no_output_when_verbose_false(capsys):
discretize('[[k_x * A(x) * k_x]]', verbose=False)
out, err = capsys.readouterr()
assert out == ''
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment