def __call__(self, equations, variables=None):
if variables is None:
variables = {}
# Get a representation of the ODE system in the form of
# dX/dt = M*X + B
varnames, matrix, constants = get_linear_system(equations)
# Make sure that the matrix M is constant, i.e. it only contains
# external variables or constant variables
t = Symbol('t', real=True, positive=True)
symbols = set.union(*(el.atoms() for el in matrix))
non_constant = _non_constant_symbols(symbols, variables, t)
if len(non_constant):
raise ValueError(('The coefficient matrix for the equations '
'contains the symbols %s, which are not '
'constant.') % str(non_constant))
# Check for time dependence
dt_var = variables.get('dt', None)
if dt_var is not None:
# This will raise an error if we meet the symbol "t" anywhere
# except as an argument of a locally constant function
for entry in itertools.chain(matrix, constants):
_check_for_locally_constant(entry, variables, dt_var.get_value(), t)
symbols = [Symbol(variable, real=True) for variable in varnames]
solution = sp.solve_linear_system(matrix.row_join(constants), *symbols)
b = sp.ImmutableMatrix([solution[symbol] for symbol in symbols]).transpose()
# Solve the system
dt = Symbol('dt', real=True, positive=True)
A = (matrix * dt).exp()
C = sp.ImmutableMatrix([A.dot(b)]) - b
_S = sp.MatrixSymbol('_S', len(varnames), 1)
updates = A * _S + C.transpose()
try:
# In sympy 0.7.3, we have to explicitly convert it to a single matrix
# In sympy 0.7.2, it is already a matrix (which doesn't have an
# is_explicit method)
updates = updates.as_explicit()
except AttributeError:
pass
# The solution contains _S[0, 0], _S[1, 0] etc. for the state variables,
# replace them with the state variable names
abstract_code = []
for idx, (variable, update) in enumerate(zip(varnames, updates)):
rhs = update
for row_idx, varname in enumerate(varnames):
rhs = rhs.subs(_S[row_idx, 0], varname)
# Do not overwrite the real state variables yet, the update step
# of other state variables might still need the original values
abstract_code.append('_' + variable + ' = ' + sympy_to_str(rhs))
# Update the state variables
for variable in varnames:
abstract_code.append('{variable} = _{variable}'.format(variable=variable))
return '\n'.join(abstract_code)
def ensure_matrix(M):
if isinstance(M, unitbridge.Quantity):
m = M.magnitude
if isinstance(m, sympy.ImmutableMatrix):
u = M.replace_magnitude(1)
return sympy.ImmutableMatrix([[i * u for i in m.row(row)]
for row in range(m.rows)])
else:
return sympy.ImmutableMatrix([M])
elif not isinstance(M, sympy.ImmutableMatrix):
return sympy.ImmutableMatrix([M])
else:
return M
def __init__(self,
R,
conjugate=False,
antisymmetry=False,
U=None,
_strict_eq=False):
if isinstance(R, sympy.ImmutableMatrix):
# If it is integer, recast to integer tinyarray
R = _make_int(R)
elif isinstance(R, ta.ndarray_int):
pass
elif isinstance(R, ta.ndarray_float):
R = _make_int(R)
elif isinstance(R, sympy.matrices.MatrixBase):
R = sympy.ImmutableMatrix(R)
R = _make_int(R)
elif isinstance(R, np.ndarray):
# If it is integer, recast to integer tinyarray
R = ta.array(R)
R = _make_int(R)
else:
raise ValueError(
'Real space rotation must be provided as a sympy matrix or an array.'
)
self.R, self.conjugate, self.antisymmetry, self.U = R, conjugate, antisymmetry, U
# Calculating sympy inverse is slow, remember it
self._strict_eq = _strict_eq
def generate_c_code(expr,syms,func_name,tmp_dir,out_dir,verbose=False,request_delete_tmp_dir=True):
tmp_dir_exists_before_call = os.path.exists(tmp_dir)
syms = list(syms)
if isinstance(expr,sympy.Matrix) or isinstance(expr,sympy.MutableMatrix) or isinstance(expr,sympy.ImmutableMatrix):
out_sym = sympy.MatrixSymbol('out_%s' % abs(hash(sympy.ImmutableMatrix(expr))), *expr.shape)
syms = syms + [out_sym]
if verbose: print "flashlight.sympy: Generating Cython code for %s..." % func_name
_autowrap_skip_compile(expr=expr,backend="cython",verbose=verbose,args=syms,tempdir=tmp_dir)
if verbose: print "flashlight.sympy: Modifying autogenerated Cython files and building..."
with open("%s/setup.py" % tmp_dir, "r") as f:
setup_py_str = f.read()
tmp_module_name = re.findall("wrapper_module_[0-9]*", setup_py_str)
tmp_autofunc_code_name = re.findall("wrapped_code_[0-9]*", setup_py_str)
assert len(tmp_module_name) == 2
assert len(tmp_autofunc_code_name) == 1
assert tmp_module_name[0] == tmp_module_name[1]
tmp_module_name = tmp_module_name[0]
tmp_autofunc_code_name = tmp_autofunc_code_name[0]
with open("%s/%s.c" % (tmp_dir,tmp_autofunc_code_name), "r") as f:
autofunc_c_str = f.read()
with open("%s/%s.h" % (tmp_dir,tmp_autofunc_code_name), "r") as f:
autofunc_h_str = f.read()
with open("%s/%s_autofunc.c" % (tmp_dir,func_name), "w") as f:
autofunc_c_str_mod = autofunc_c_str
autofunc_c_str_mod = autofunc_c_str_mod.replace("autofunc(", "%s_autofunc(" % func_name)
autofunc_c_str_mod = autofunc_c_str_mod.replace('#include "%s.h"' % tmp_autofunc_code_name, '#include "%s.h"' % (func_name+"_autofunc"))
f.write(autofunc_c_str_mod)
with open("%s/%s_autofunc.h" % (tmp_dir,func_name), "w") as f:
autofunc_h_str_mod = autofunc_h_str
autofunc_h_str_mod = autofunc_h_str.replace("autofunc(", "%s_autofunc(" % func_name)
autofunc_h_str_mod = autofunc_h_str_mod.replace("AUTOWRAP__%s__H" % tmp_autofunc_code_name.upper(), "%s_H" % (func_name+"_autofunc").upper())
f.write(autofunc_h_str_mod)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
cmd = "cp %s/%s_autofunc.c %s/%s_autofunc.c" % (tmp_dir,func_name,out_dir,func_name)
if verbose: print cmd
output = subprocess.check_output(cmd, shell=True)
if verbose and len(output) > 0: print output
cmd = "cp %s/%s_autofunc.h %s/%s_autofunc.h" % (tmp_dir,func_name,out_dir,func_name)
if verbose: print cmd
output = subprocess.check_output(cmd, shell=True)
if verbose and len(output) > 0: print output
if not tmp_dir_exists_before_call and request_delete_tmp_dir:
cmd = "rm -rf %s" % tmp_dir
if verbose: print cmd
output = subprocess.check_output(cmd, shell=True)
if verbose and len(output) > 0: print output
if verbose: print
def hess(self):
""" `~sympy.matrices.immutable.ImmutableDenseMatrix` : \
Second derivatives of :py:attr:`expr` with respect to \
:py:attr:`prob.vars`. """
if self._hess is None:
self._hess = sym.ImmutableMatrix(
self.grad.jacobian(self.prob.var_scalars))
return self._hess
def grad(self):
""" `~sympy.matrices.immutable.ImmutableDenseMatrix` : \
First derivatives of :py:attr:`expr` with respect to \
:py:attr:`prob.vars`. """
if self._grad is None:
self._grad = sym.ImmutableMatrix(
[self.expr.diff(s) for s in self.prob.var_scalars])
return self._grad
def update_subs(self, subs):
if isinstance(self.length, math.Symbol):
self.length = self.length.subs(subs)
if self.pts_2D is not None and len(self.pts_2D) == 2 and len(
self.pts_2D[0]) == 2:
self.pts_2D = tuple(
sympy.ImmutableMatrix([
dim.subs(subs) if isinstance(dim, sympy.Basic) else dim
for dim in p
]) for p in self.pts_2D)
if self.pts_3D is not None and len(self.pts_3D) == 2 and len(
self.pts_3D[0]) == 3:
self.pts_3D = tuple(
sympy.ImmutableMatrix([
dim.subs(subs) if isinstance(dim, sympy.Basic) else dim
for dim in p
]) for p in self.pts_3D)
def __init__(self, alpha, delta):
H = sp.ImmutableMatrix(
[[1 / 2 * alpha * _t, 1 / 2 * delta],
[1 / 2 * sp.conjugate(delta), -1 / 2 * alpha * _t]])
if not sp.S(delta).has(_t):
self._lz_value = self.lz_formula(alpha, delta)
else:
self._lz_value = None
QuantumSystem.__init__(self, H)
def _mul(R1, R2):
# Cached multiplication of spatial parts.
if is_sympy_matrix(R1) and is_sympy_matrix(R2):
# If spatial parts are sympy matrices, use cached multiplication.
R = R1 * R2
elif not (is_sympy_matrix(R1) or is_sympy_matrix(R2)):
# If arrays, use dot
R = ta.dot(R1, R2)
elif ((is_sympy_matrix(R1) or is_sympy_matrix(R2)) and
(isinstance(R1, ta.ndarray_int) or isinstance(R2, ta.ndarray_int))):
# Multiplying sympy and integer tinyarray is ok, should result in sympy
R = sympy.ImmutableMatrix(R1) * sympy.ImmutableMatrix(R2)
else:
raise ValueError("Mixing of sympy and floating point in the spatial part R is not allowed. "
"To avoid this error, make sure that all PointGroupElements are initialized "
"with either floating point arrays or sympy matrices as rotations. "
"Integer arrays are allowed in both cases.")
R = _make_int(R)
return R
def UnaryFunc(self, node, **kw):
arg = self(node.arg)
op = self(node.op)
if isinstance(arg, sp.MatrixBase) and isinstance(node.op, IsElemwise):
if arg.shape[1] != 1:
raise NotImplementedError(
"Elementwise operations are only supported for vectors (column matrices)"
)
return sp.ImmutableMatrix([op(arg[i]) for i in range(len(arg))])
else:
return op(arg)
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 got == out,\
"Should be: _split_factors({})=={}. Not {}".format(inp, out, got)
def get_foreman():
# Basis transformation
v_n = lambda i, N: sympy.Matrix( # noqa: E731
[1 if i == n else 0 for n in range(N)])
S, X, Y, Z = [v_n(i, 4) for i in range(4)]
up, dw = [v_n(i, 2) for i in range(2)]
X = sympy.I * X
Y = sympy.I * Y
Z = sympy.I * Z
molenkamp_basis = [
kr(up, S),
kr(dw, S),
+(1 / sympy.sqrt(2)) * kr(up, X + sympy.I * Y),
+(1 / sympy.sqrt(6)) * (kr(dw, X + sympy.I * Y) - kr(up, 2 * Z)),
-(1 / sympy.sqrt(6)) * (kr(up, X - sympy.I * Y) + kr(dw, 2 * Z)),
-(1 / sympy.sqrt(2)) * kr(dw, X - sympy.I * Y),
+(1 / sympy.sqrt(3)) * (kr(dw, X + sympy.I * Y) + kr(up, Z)),
+(1 / sympy.sqrt(3)) * (kr(up, X - sympy.I * Y) - kr(dw, Z)),
]
subs_notation = {
Ec: Ev + E0,
L: -(g1 + 4 * g2) * (hbar**2 / 2 / m0),
M: -(g1 - 2 * g2) * (hbar**2 / 2 / m0),
Np: -(3 * g3 + (3 * kappa + 1)) * (hbar**2 / 2 / m0),
Nm: -(3 * g3 - (3 * kappa + 1)) * (hbar**2 / 2 / m0),
Ac: (g0 * hbar**2 / 2 / m0),
}
Hs = spin_orbit()
Hcc = sympy.Matrix([Ec + kx * Ac * kx + ky * Ac * ky + kz * Ac * kz])
Hcv = +sympy.I * sympy.Matrix([[P * kx, P * ky, P * kz]])
Hvc = -sympy.I * sympy.Matrix([kx * P, ky * P, kz * P])
data = [[valence_term(i, j) for j in range(3)] for i in range(3)]
Hvv = sympy.Matrix(data)
H4 = sympy.BlockMatrix([[Hcc, Hcv], [Hvc, Hvv]])
H8 = sympy.Matrix(sympy.BlockDiagMatrix(H4, H4)) + Hs
dag = lambda x: x.conjugate().transpose() # noqa: E731
U_dag = sympy.Matrix(sympy.BlockMatrix([molenkamp_basis]))
U = dag(U_dag)
hamiltonian = U * H8 * dag(U)
hamiltonian = hamiltonian.subs(subs_notation)
hamiltonian = hamiltonian + (Ev - Delta / 3) * sympy.diag(
0, 0, 1, 1, 1, 1, 1, 1)
return sympy.ImmutableMatrix(hamiltonian.expand())
def gen_sigmas():
"""generate Pauli matrices.
Usage:\n
sp.Matrix(np.dot(vec, sigmas).reshape(2,2))
"""
sigmas = []
for i in range(4):
if (i == 0):
sigma = sp.Identity(2).as_explicit()
else:
sigma = sp.ImmutableMatrix(msigma(i))
sigmas.append(sigma)
return sigmas
def BinOp(self, node, **kw):
left = self(node.left)
right = self(node.right)
op = self(node.op)
if isinstance(left, sp.MatrixBase) and isinstance(
right, sp.MatrixBase) and isinstance(node.op, IsElemwise):
if left.shape != right.shape:
raise ValueError("Shapes of the arguments do not match")
if left.shape[1] != 1:
raise NotImplementedError(
"Elementwise operations are only supported for vectors (column matrices)"
)
return sp.ImmutableMatrix(
[op(left[i], right[i]) for i in range(len(left))])
else:
return op(left, right)
def _process_bound(self, b):
if self.is_scalar:
if not util.is_scalar(b):
raise TypeError(
f"Variable {self.symbol} is a scalar but non-scalar bounds"
f"were provided.")
else:
shape = self.symbol.shape
try:
b = sym.ImmutableMatrix(b).reshape(*shape)
except (TypeError, ValueError):
b = b * sym.ones(*shape)
b = sympify(b)
if not self.prob.depends_only_on_params(b):
raise ValueError(
"Bounds can depend only on the symbols in prob.params")
return b
def IfThenElse(self, node):
# Piecewise function with matrix output is not a Matrix itself, which breaks some of the logic
# Hence this won't work in general:
test, iftrue, iffalse = self(node.test), self(node.iftrue), self(
node.iffalse)
if hasattr(iftrue, 'shape'):
if iftrue.shape != iffalse.shape:
raise ValueError("Shapes of the IF branches must match")
if iftrue.shape[1] != 1:
raise NotImplementedError(
"Elementwise operations are only supported for vectors (column matrices)"
)
return sp.ImmutableMatrix([
sp.Piecewise((ift, test), (iff, True))
for ift, iff in zip(iftrue, iffalse)
])
else:
return sp.Piecewise((iftrue, test), (iffalse, True))
def test_spatial_types():
S1 = PointGroupElement(sympy.eye(2), False, False, np.eye(3))
S2 = PointGroupElement(sympy.Matrix([[0, 1], [1, 0]]), True, False,
np.eye(3))
S3 = PointGroupElement(np.eye(2), False, False, 1j * np.eye(3))
C6s = PointGroupElement(
sympy.ImmutableMatrix([[sympy.Rational(1, 2),
sympy.sqrt(3) / 2],
[-sympy.sqrt(3) / 2,
sympy.Rational(1, 2)]]))
C6f = PointGroupElement(
np.array([[1 / 2, np.sqrt(3) / 2], [-np.sqrt(3) / 2, 1 / 2]]))
assert S2**2 == S1
assert not S1 == S2
assert S1 == S3
assert C6s == C6f
# Mixing sympy with other types raises an error
with raises(ValueError):
S = C6s * C6f
def __init__(self,
symbolic_expression: sym.Expr,
argument_symbols: tp.Tuple[sym.Symbol, ...],
time_symbol: tp.Optional[sym.Symbol] = None,
symbolic_time_function_name_to_numeric_time_function_map: tp.
Optional[tp.Dict[str, tp.Callable]] = None,
perform_cse: bool = True,
lazy_initialization: bool = False,
pre_attach: bool = True,
dtype: np.generic = np.float32):
super().__init__(
argument_symbols=argument_symbols,
time_symbol=time_symbol,
symbolic_vector_expression=sym.ImmutableMatrix(
[symbolic_expression]),
symbolic_time_function_name_to_numeric_time_function_map=
symbolic_time_function_name_to_numeric_time_function_map,
perform_cse=perform_cse,
lazy_initialization=lazy_initialization,
pre_attach=pre_attach,
dtype=dtype)
def __init__(self,
R,
conjugate=False,
antisymmetry=False,
U=None,
_strict_eq=False,
*,
locals=None):
if isinstance(R, sympy.ImmutableMatrix):
# If it is integer, recast to integer tinyarray
R = _make_int(R)
elif isinstance(R, ta.ndarray_int):
pass
elif isinstance(R, ta.ndarray_float):
R = _make_int(R)
elif isinstance(R, MatrixBase):
R = sympy.ImmutableMatrix(R)
R = _make_int(R)
elif isinstance(R, np.ndarray):
# If it is integer, recast to integer tinyarray
R = ta.array(R)
R = _make_int(R)
else:
raise ValueError(
'Real space rotation must be provided as a sympy matrix or an array.'
)
# Normalize U
if U is None:
pass
else:
try:
U = np.atleast_2d(np.array(U, dtype=complex))
except (ValueError, TypeError):
U = sympify(U, locals=locals)
U = np.atleast_2d(np.array(U, dtype=complex))
self.R, self.conjugate, self.antisymmetry, self.U = R, conjugate, antisymmetry, U
# Calculating sympy inverse is slow, remember it
self._strict_eq = _strict_eq
def _build_hamiltonian(self):
# return foreman(self._parameter_coords, self.components, self.bands)
if self._parameter_coords is not None:
self._parameter_coords = validate_coords(self._parameter_coords)
str_coords = "({})".format(", ".join(self._parameter_coords))
subs = {v: v + str_coords for v in self._varied_parameters}
else:
subs = {}
hamiltonian_components = [
kwant.continuum.sympify(_models_cache[c], locals=subs)
for c in self.components
]
hamiltonian = sympy.ImmutableMatrix(
sympy.MatAdd(*hamiltonian_components))
indices = []
for band in self.bands:
indices += self._band_indices[band]
return hamiltonian[:, indices][indices, :]
import sympy as sp
class Permutation(tuple):
def __mul__(self, s):
return Permutation(self[x] for x in s)
def sg(elts):
a = {}
b = set(elts)
while a != b:
a = set(b)
b |= {s1 * s2 for s1 in b for s2 in b}
return b
print(sg({Permutation(x) for x in ((1, 2, 3, 0), (3, 2, 1, 0))}))
print(
sg({sp.ImmutableMatrix(x)
for x in (((1, 0), (0, -1)), ((0, -1), (1, 0)))}))
def __init__(self, alpha, delta):
H = sp.ImmutableMatrix([[alpha * _t, delta, 0],
[delta.conjugate(), 0, delta],
[0, delta.conjugate(), -alpha * _t]])
QuantumSystem.__init__(self, H)
def test_bloch():
# Simple tests for Bloch models
# Hexagonal point group
eyesym = sympy.ImmutableMatrix(sympy.eye(2))
Mx = PointGroupElement(sympy.ImmutableMatrix([[-1, 0], [0, 1]]), False,
False, None)
C6 = PointGroupElement(
sympy.ImmutableMatrix([[sympy.Rational(1, 2),
sympy.sqrt(3) / 2],
[-sympy.sqrt(3) / 2,
sympy.Rational(1, 2)]]), False, False, None)
TR = time_reversal(realspace_dim=2)
PH = particle_hole(realspace_dim=2)
gens_hex_2D = {Mx, C6, TR, PH}
hex_group_2D = generate_group(gens_hex_2D)
assert len(hex_group_2D) == 48
# First example
ham6 = 'm * (cos(k_x) + cos(1/2*k_x + sqrt(3)/2*k_y) + cos(-1/2*k_x + sqrt(3)/2*k_y))'
H6 = Model(ham6, momenta=['k_x', 'k_y'])
sg, Ps = discrete_symmetries(H6, hex_group_2D)
assert [P.shape for P in Ps] == [(1, 1, 1)]
assert len(sg) == 24
assert sg == generate_group({Mx, C6, TR})
# extend model to add SOC
ham62 = 'eye(2) * (' + ham6 + ') +'
ham62 += 'alpha * (sin(k_x) * sigma_x + sin(1/2*k_x + sqrt(3)/2*k_y) * (1/2*sigma_x + sqrt(3)/2*sigma_y) +'
ham62 += 'sin(-1/2*k_x + sqrt(3)/2*k_y) * (-1/2*sigma_x + sqrt(3)/2*sigma_y))'
H62 = Model(ham62, momenta=['k_x', 'k_y'])
sg, Ps = discrete_symmetries(H62, hex_group_2D)
assert [P.shape for P in Ps] == [(1, 2, 2)]
assert len(sg) == 24
assert sg == generate_group({Mx, C6, TR})
# Add degeneracy
ham63 = 'kron(eye(2), ' + ham62 + ')'
H63 = Model(ham63, momenta=['k_x', 'k_y'])
sg, Ps = discrete_symmetries(H63, hex_group_2D)
assert [P.shape for P in Ps] == [(2, 4, 2)]
assert len(sg) == 24
assert sg == generate_group({Mx, C6, TR})
# Add PH states
ham64 = 'kron(sigma_z, ' + ham62 + ')'
H64 = Model(ham64, momenta=['k_x', 'k_y'])
sg, Ps = discrete_symmetries(H64, hex_group_2D)
assert [P.shape for P in Ps] == [(1, 4, 2), (1, 4, 2)]
assert len(sg) == 48
assert sg == generate_group({Mx, C6, TR, PH})
# Test sparse
H64 = Model(ham64, momenta=['k_x', 'k_y'])
sg, Ps = discrete_symmetries(H64, hex_group_2D, sparse_linalg=True)
assert [P.shape for P in Ps] == [(1, 4, 2), (1, 4, 2)]
assert len(sg) == 48
assert sg == generate_group({Mx, C6, TR, PH})
Hcsr = H64.tocsr()
sg, Ps = discrete_symmetries(Hcsr, hex_group_2D, sparse_linalg=True)
assert [P.shape for P in Ps] == [(1, 4, 2), (1, 4, 2)]
assert len(sg) == 48
assert sg == generate_group({Mx, C6, TR, PH})
sg, Ps = discrete_symmetries(Hcsr, hex_group_2D, sparse_linalg=False)
assert [P.shape for P in Ps] == [(1, 4, 2), (1, 4, 2)]
assert len(sg) == 48
assert sg == generate_group({Mx, C6, TR, PH})
def __init__(self, alpha, delta, epsilon):
H = sp.ImmutableMatrix([[alpha * _t + epsilon, delta, 0],
[delta.conjugate(), 0, delta],
[0, delta.conjugate(), alpha * _t - epsilon]])
QuantumSystem.__init__(self, H)
def __call__(self, equations, variables=None, method_options=None):
method_options = extract_method_options(method_options,
{'simplify': True})
if equations.is_stochastic:
raise UnsupportedEquationsException("Cannot solve stochastic "
"equations with this state "
"updater.")
if variables is None:
variables = {}
# Get a representation of the ODE system in the form of
# dX/dt = M*X + B
varnames, matrix, constants = get_linear_system(equations, variables)
# No differential equations, nothing to do (this occurs sometimes in the
# test suite where the whole model is nothing more than something like
# 'v : 1')
if matrix.shape == (0, 0):
return ''
# Make sure that the matrix M is constant, i.e. it only contains
# external variables or constant variables
# Check for time dependence
dt_value = variables['dt'].get_value(
)[0] if 'dt' in variables else None
# This will raise an error if we meet the symbol "t" anywhere
# except as an argument of a locally constant function
for entry in itertools.chain(matrix, constants):
if not is_constant_over_dt(entry, variables, dt_value):
raise UnsupportedEquationsException(
f"Expression '{sympy_to_str(entry)}' is not guaranteed to be "
f"constant over a time step.")
symbols = [Symbol(variable, real=True) for variable in varnames]
solution = sp.solve_linear_system(matrix.row_join(constants), *symbols)
if solution is None or set(symbols) != set(solution.keys()):
raise UnsupportedEquationsException("Cannot solve the given "
"equations with this "
"stateupdater.")
b = sp.ImmutableMatrix([solution[symbol] for symbol in symbols])
# Solve the system
dt = Symbol('dt', real=True, positive=True)
try:
A = (matrix * dt).exp()
except NotImplementedError:
raise UnsupportedEquationsException("Cannot solve the given "
"equations with this "
"stateupdater.")
if method_options['simplify']:
A = A.applyfunc(
lambda x: sp.factor_terms(sp.cancel(sp.signsimp(x))))
C = sp.ImmutableMatrix(A * b) - b
_S = sp.MatrixSymbol('_S', len(varnames), 1)
updates = A * _S + C
updates = updates.as_explicit()
# The solution contains _S[0, 0], _S[1, 0] etc. for the state variables,
# replace them with the state variable names
abstract_code = []
for idx, (variable, update) in enumerate(zip(varnames, updates)):
rhs = update
if rhs.has(I, re, im):
raise UnsupportedEquationsException(
"The solution to the linear system "
"contains complex values "
"which is currently not implemented.")
for row_idx, varname in enumerate(varnames):
rhs = rhs.subs(_S[row_idx, 0], varname)
# Do not overwrite the real state variables yet, the update step
# of other state variables might still need the original values
abstract_code.append(f"_{variable} = {sympy_to_str(rhs)}")
# Update the state variables
for variable in varnames:
abstract_code.append(f"{variable} = _{variable}")
return '\n'.join(abstract_code)
def build_module_autowrap(expr,
syms,
module_name,
tmp_dir,
out_dir,
dummify=False,
cse=False,
cse_ordering="canonical",
build_vectorized=False,
pretty_dummy_symbol_names=False,
verbose=False,
request_delete_tmp_dir=True):
tmp_dir_exists_before_call = os.path.exists(tmp_dir)
is_matrix_expr = isinstance(expr, sympy.Matrix) or isinstance(
expr, sympy.MutableMatrix) or isinstance(expr, sympy.ImmutableMatrix)
if dummify:
if verbose:
print "flashlight.sympy: Generating dummy symbols for %s..." % module_name
expr, syms = _dummify(expr, syms, pretty_dummy_symbol_names)
if verbose:
print "flashlight.sympy: Finished generating dummy symbols for %s." % module_name
if cse:
include_str, subexpr_eval_str = generate_c_code_cse(
expr,
syms,
module_name,
tmp_dir,
out_dir,
cse_ordering=cse_ordering,
verbose=verbose,
request_delete_tmp_dir=request_delete_tmp_dir)
if is_matrix_expr:
expr = sympy.Matrix.zeros(expr.rows, expr.cols)
else:
expr = 0
syms = list(syms)
if is_matrix_expr:
out_sym = sympy.MatrixSymbol(
'out_%s' % abs(hash(sympy.ImmutableMatrix(expr))), *expr.shape)
syms = syms + [out_sym]
if verbose:
print "flashlight.sympy: Generating Cython code for %s..." % module_name
_autowrap_skip_compile(expr=expr,
backend="cython",
verbose=verbose,
args=syms,
tempdir=tmp_dir)
if verbose:
print "flashlight.sympy: Modifying autogenerated Cython files for %s" % module_name
with open("%s/setup.py" % tmp_dir, "r") as f:
setup_py_str = f.read()
tmp_module_name = re.findall("wrapper_module_[0-9]*", setup_py_str)
tmp_autofunc_code_name = re.findall("wrapped_code_[0-9]*",
setup_py_str)
assert len(tmp_module_name) == 2
assert len(tmp_autofunc_code_name) == 1
assert tmp_module_name[0] == tmp_module_name[1]
tmp_module_name = tmp_module_name[0]
tmp_autofunc_code_name = tmp_autofunc_code_name[0]
with open("%s/%s.pyx" % (tmp_dir, tmp_module_name), "r") as f:
wrapper_module_pyx_str = f.read()
with open("%s/%s.c" % (tmp_dir, tmp_autofunc_code_name), "r") as f:
autofunc_c_str = f.read()
with open("%s/%s.h" % (tmp_dir, tmp_autofunc_code_name), "r") as f:
autofunc_h_str = f.read()
with open("%s/%s_setup.py" % (tmp_dir, module_name), "w") as f:
setup_py_str_mod = setup_py_str
setup_py_str_mod = setup_py_str_mod.replace(
"from Cython.Distutils import build_ext",
"from Cython.Distutils import build_ext\n\nimport numpy\nimport os\n\nos.environ['CC'] = '%s'\nos.environ['CXX'] = '%s'\n\n"
% (cc_compiler, cxx_compiler))
setup_py_str_mod = setup_py_str_mod.replace(
"extra_compile_args=['-std=c99']",
"extra_compile_args=['-std=c99','-fno-var-tracking','-fno-var-tracking-assignments'], include_dirs=[numpy.get_include()]"
)
setup_py_str_mod = setup_py_str_mod.replace("%s" % tmp_module_name,
"%s" % module_name)
setup_py_str_mod = setup_py_str_mod.replace(
"%s" % tmp_autofunc_code_name, "%s_autofunc" % module_name)
f.write(setup_py_str_mod)
with open("%s/%s.pyx" % (tmp_dir, module_name), "w") as f:
wrapper_module_pyx_str_mod = wrapper_module_pyx_str
wrapper_module_pyx_str_mod = wrapper_module_pyx_str_mod.replace(
"%s" % tmp_autofunc_code_name, "%s_autofunc" % module_name)
f.write(wrapper_module_pyx_str_mod)
if not cse:
with open("%s/%s_autofunc.c" % (tmp_dir, module_name), "w") as f:
autofunc_c_str_mod = autofunc_c_str
autofunc_c_str_mod = autofunc_c_str_mod.replace(
'#include "%s.h"' % tmp_autofunc_code_name,
'#include "%s.h"' % (module_name + "_autofunc"))
f.write(autofunc_c_str_mod)
with open("%s/%s_autofunc.h" % (tmp_dir, module_name), "w") as f:
autofunc_h_str_mod = autofunc_h_str
autofunc_h_str_mod = autofunc_h_str_mod.replace(
"AUTOWRAP__%s__H" % tmp_autofunc_code_name.upper(),
"%s_H" % (module_name + "_autofunc").upper())
f.write(autofunc_h_str_mod)
if build_vectorized:
if verbose:
print "flashlight.sympy: Generating Cython code for %s_vectorized..." % module_name
if is_matrix_expr:
c_return_type_str = "void"
c_out_var_str = "out_123456789"
c_func_signature_str = str([
"double %s" % str(sym) for sym in syms[:-1]
])[1:-1].replace("'", "") + ", double *%s" % syms[-1]
cython_out_type_str = "numpy.ndarray[FLOAT64_DTYPE_t, ndim=3]"
cython_out_shape_str = "(args.shape[0],%s,%s)" % (expr.rows,
expr.cols)
cython_out_v_type_str = "FLOAT64_DTYPE_t [:,:,:]"
cython_args_str = str([
"args[i,%d]" % i for i in range(len(syms[:-1]))
])[1:-1].replace("'", "") + ", &out_v[i,0,0]"
cython_loop_body_str = "autofunc(%s)" % cython_args_str
else:
c_return_type_str = "double"
c_func_signature_str = str([
"double %s" % str(sym) for sym in syms
])[1:-1].replace("'", "")
cython_out_type_str = "numpy.ndarray[FLOAT64_DTYPE_t, ndim=1]"
cython_out_shape_str = "(args.shape[0])"
cython_out_v_type_str = "FLOAT64_DTYPE_t [:]"
cython_args_str = str(["args[i,%d]" % i for i in range(len(syms))
])[1:-1].replace("'", "")
cython_loop_body_str = "out_v[i] = autofunc(%s)" % cython_args_str
vectorized_setup_py_str_eval = vectorized_setup_py_str % (
cc_compiler, cxx_compiler, module_name, module_name, module_name)
vectorized_pyx_str_eval = vectorized_pyx_str % (
module_name, c_return_type_str, c_func_signature_str,
cython_out_type_str, cython_out_shape_str, cython_out_v_type_str,
cython_loop_body_str)
with open("%s/%s_vectorized_setup.py" % (tmp_dir, module_name),
"w") as f:
vectorized_setup_py_str_mod = vectorized_setup_py_str_eval
f.write(vectorized_setup_py_str_mod)
with open("%s/%s_vectorized.pyx" % (tmp_dir, module_name), "w") as f:
vectorized_pyx_str_mod = vectorized_pyx_str_eval
f.write(vectorized_pyx_str_mod)
cwd = os.getcwd()
try:
os.chdir(tmp_dir)
if verbose:
print "flashlight.sympy: Building Cython code for %s..." % module_name
cmd = "python %s_setup.py build_ext --inplace" % module_name
if verbose: print cmd
output = subprocess.check_output(cmd, shell=True)
if verbose and len(output) > 0: print output
if build_vectorized:
if verbose:
print "flashlight.sympy: Building Cython code for %s_vectorized..." % module_name
cmd = "python %s_vectorized_setup.py build_ext --inplace" % module_name
if verbose: print cmd
output = subprocess.check_output(cmd, shell=True)
if verbose and len(output) > 0: print output
finally:
os.chdir(cwd)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
cmd = "cp %s/%s.so %s/%s.so" % (tmp_dir, module_name, out_dir, module_name)
if verbose: print cmd
output = subprocess.check_output(cmd, shell=True)
if verbose and len(output) > 0: print output
if build_vectorized:
cmd = "cp %s/%s_vectorized.so %s/%s_vectorized.so" % (
tmp_dir, module_name, out_dir, module_name)
if verbose: print cmd
output = subprocess.check_output(cmd, shell=True)
if verbose and len(output) > 0: print output
# cmd = "cp %s/%s_autofunc.c %s/%s_autofunc.c" % (tmp_dir,module_name,out_dir,module_name)
# if verbose: print cmd
# output = subprocess.check_output(cmd, shell=True)
# if verbose and len(output) > 0: print output
# cmd = "cp %s/%s_autofunc.h %s/%s_autofunc.h" % (tmp_dir,module_name,out_dir,module_name)
# if verbose: print cmd
# output = subprocess.check_output(cmd, shell=True)
# if verbose and len(output) > 0: print output
if not tmp_dir_exists_before_call and request_delete_tmp_dir:
cmd = "rm -rf %s" % tmp_dir
if verbose: print cmd
output = subprocess.check_output(cmd, shell=True)
if verbose and len(output) > 0: print output
def __init__(self, alpha, delta, epsilon):
H = sp.ImmutableMatrix([[epsilon + alpha * _t, delta, 0],
[delta.conjugate(), -2 * epsilon, delta],
[0, delta.conjugate(), epsilon - alpha * _t]])
QuantumSystem.__init__(self, H)
def __call__(self, equations, variables=None, simplify=True):
if equations.is_stochastic:
raise UnsupportedEquationsException('Cannot solve stochastic '
'equations with this state '
'updater.')
if variables is None:
variables = {}
# Get a representation of the ODE system in the form of
# dX/dt = M*X + B
varnames, matrix, constants = get_linear_system(equations, variables)
# No differential equations, nothing to do (this occurs sometimes in the
# test suite where the whole model is nothing more than something like
# 'v : 1')
if matrix.shape == (0, 0):
return ''
# Make sure that the matrix M is constant, i.e. it only contains
# external variables or constant variables
t = Symbol('t', real=True, positive=True)
# Check for time dependence
if 'dt' in variables:
dt_value = variables['dt'].get_value()[0]
# This will raise an error if we meet the symbol "t" anywhere
# except as an argument of a locally constant function
t = Symbol('t', real=True, positive=True)
for entry in itertools.chain(matrix, constants):
_check_for_locally_constant(entry, variables, dt_value, t)
symbols = [Symbol(variable, real=True) for variable in varnames]
solution = sp.solve_linear_system(matrix.row_join(constants), *symbols)
if solution is None or set(symbols) != set(solution.keys()):
raise UnsupportedEquationsException('Cannot solve the given '
'equations with this '
'stateupdater.')
b = sp.ImmutableMatrix([solution[symbol]
for symbol in symbols]).transpose()
# Solve the system
dt = Symbol('dt', real=True, positive=True)
try:
A = (matrix * dt).exp()
except NotImplementedError:
raise UnsupportedEquationsException('Cannot solve the given '
'equations with this '
'stateupdater.')
if simplify:
A = A.applyfunc(
lambda x: sp.factor_terms(sp.cancel(sp.signsimp(x))))
C = sp.ImmutableMatrix([A.dot(b)]) - b
_S = sp.MatrixSymbol('_S', len(varnames), 1)
updates = A * _S + C.transpose()
updates = updates.as_explicit()
# The solution contains _S[0, 0], _S[1, 0] etc. for the state variables,
# replace them with the state variable names
abstract_code = []
for idx, (variable, update) in enumerate(zip(varnames, updates)):
rhs = update
for row_idx, varname in enumerate(varnames):
rhs = rhs.subs(_S[row_idx, 0], varname)
# Do not overwrite the real state variables yet, the update step
# of other state variables might still need the original values
abstract_code.append('_' + variable + ' = ' + sympy_to_str(rhs))
# Update the state variables
for variable in varnames:
abstract_code.append(
'{variable} = _{variable}'.format(variable=variable))
return '\n'.join(abstract_code)
def hexagonal(dim=2,
tr=True,
ph=True,
generators=False,
sympy_R=True,
spin=None):
"""
Generate hexagonal point group in standard basis in 2 or 3 dimensions.
Mirror symmetries with the main coordinate axes as normals are included.
In 3D the hexagonal axis is the z axis.
Parameters:
-----------
dim : int (default 2)
Real sapce dimensionality, 2 or 3.
tr, ph : bool (default True)
Whether to include time-reversal and particle-hole
symmetry.
generators : bool (default True)
Only return the group generators if True.
sympy_R: bool (default True)
Whether the rotation matrices should be exact sympy
representations.
spin : float or sequence of arrays (optional)
Spin representation to use for the unitary action of the
operator. If not provided, the PointGroupElements have the unitary
action set to None. If float is provided, it should be integer or half-integer
specifying the spin representation in the standard basis, see `spin_matrices`.
Otherwise a sequence of 3 arrays of identical square size must be provided
representing 3 components of the angular momentum operator. The unitary action
of rotation operator is `U = exp(-i n⋅s)`. In 2D the z axis is assumed to be
the rotation axis. If `ph` is True, `spin` may not be provided, as it is not
possible to deduce the unitary representation of particle-hole symmetry from
spin alone. In this case construct the particle-hole operator manually.
Returns:
--------
set of PointGroupElements
Notes:
------
Warning: in 2D the real space action of a mirror and and a 2-fold rotation
around an axis in the plane is identical, however the action on angular momentum
is different. Here we consider the action of the mirror, which is the same as the
action of a 2-fold rotation around the mirror axis, assuming inversion acts trivially.
"""
if spin is not None and sympy_R:
U6 = spin_rotation(np.pi / 3 * np.array([0, 0, 1]), spin)
else:
U6 = None
if dim == 2:
Mx = mirror([1, 0], spin=spin)
if sympy_R:
C6 = PointGroupElement(
sympy.ImmutableMatrix(
[[sympy.Rational(1, 2),
sympy.sqrt(3) / 2],
[-sympy.sqrt(3) / 2,
sympy.Rational(1, 2)]]), False, False, U6)
else:
C6 = rotation(1 / 6, spin=spin)
gens = {Mx, C6}
elif dim == 3:
I = inversion(
3, U=(None if spin is None else spin_rotation(np.zeros(3), spin)))
C2x = rotation(1 / 2, [1, 0, 0], spin=spin)
if sympy_R:
C6 = PointGroupElement(
sympy.ImmutableMatrix(
[[sympy.Rational(1, 2),
sympy.sqrt(3) / 2, 0],
[-sympy.sqrt(3) / 2,
sympy.Rational(1, 2), 0], [0, 0, 1]]), False, False, U6)
else:
C6 = rotation(1 / 6, [0, 0, 1], spin=spin)
gens = {I, C2x, C6}
else:
raise ValueError('Only 2 and 3 dimensions are supported.')
if tr:
TR = time_reversal(dim, spin=spin)
gens.add(TR)
if ph:
PH = particle_hole(dim)
gens.add(PH)
if generators:
return gens
else:
return generate_group(gens)
def __init__(self,
symbolic_array_expressions: tp.Tuple[tp.Union[sym.Array,
sym.MatrixBase],
...],
argument_symbols: tp.Tuple[sym.Symbol],
time_symbol: tp.Optional[sym.Symbol] = None,
symbolic_time_function_name_to_numeric_time_function_map: tp.
Optional[tp.Dict[str, tp.Callable]] = None,
squeeze_column_vectors: tp.Optional[tp.Tuple[bool,
...]] = None,
perform_cse: bool = True,
lazy_initialization: bool = False,
pre_attach: tp.Optional[tp.Tuple[bool, ...]] = None,
dtype: np.generic = np.float32):
if time_symbol is None:
if symbolic_time_function_name_to_numeric_time_function_map is not None and \
len(symbolic_time_function_name_to_numeric_time_function_map) != 0:
raise ValueError(
'Argument time_symbol must not be none if '
'symbolic_time_function_name_to_numeric_time_function_map '
'is provided.')
if any([
DUMMY_TIME_SYMBOL in symbolic_array_expression.free_symbols
for symbolic_array_expression in symbolic_array_expressions
]):
raise ValueError('Argument time_symbol must be provided '
'explicity if expression contains '
'DUMMY_TIME_SYMBOL.')
time_symbol = DUMMY_TIME_SYMBOL
self._time_symbol_provided = False
else:
self._time_symbol_provided = True
if symbolic_time_function_name_to_numeric_time_function_map is None:
symbolic_time_function_name_to_numeric_time_function_map = dict()
symbolic_array_expressions \
= tuple(check_and_make_sequence(symbolic_array_expressions,
sym.MatrixBase))
# print(f'{symbolic_array_expressions=}\n')
symbolic_array_expressions = \
tuple(sym.ImmutableMatrix(symbolic_array_expression)
if isinstance(symbolic_array_expression, sym.MatrixBase)
else sym.ImmutableDenseNDimArray(symbolic_array_expression)
for symbolic_array_expression
in symbolic_array_expressions)
if pre_attach is None:
pre_attach = len(symbolic_array_expressions) * [True]
pre_attach = check_and_make_sequence(pre_attach, bool)
if squeeze_column_vectors is None:
squeeze_column_vectors = len(symbolic_array_expressions) * [False]
squeeze_column_vectors = check_and_make_sequence(
squeeze_column_vectors, bool)
if not (len(symbolic_array_expressions) == len(pre_attach)
and len(symbolic_array_expressions)
== len(squeeze_column_vectors)):
raise ValueError(
"symbolic_matrix_expressions, pre_attach, and "
"squeeze_column_vectors must have the same number "
"of elements")
self._argument_symbols = argument_symbols
self._time_symbol = time_symbol
self._symbolic_array_expressions = symbolic_array_expressions
self._symbolic_time_function_name_to_numeric_time_function_map = \
symbolic_time_function_name_to_numeric_time_function_map
self._squeeze_column_vectors = squeeze_column_vectors
self._symbols = \
tuple([time_symbol] +
list(argument_symbols))
self._perform_cse = perform_cse
self._pre_attach = pre_attach
self._dtype = dtype
self._shapes \
= tuple([symbolic_array_expression.shape
for symbolic_array_expression
in symbolic_array_expressions])
if lazy_initialization:
self._functions_cpu = None
self._functions_gpu = None
else:
self._perform_initialization()
