Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# Copyright 2011-2017 Kwant authors. 

2# 

3# This file is part of Kwant. It is subject to the license terms in the file 

4# LICENSE.rst found in the top-level directory of this distribution and at 

5# https://kwant-project.org/license. A list of Kwant authors can be found in 

6# the file AUTHORS.rst at the top-level directory of this distribution and at 

7# https://kwant-project.org/authors. 

8 

9import keyword 

10from collections import defaultdict 

11 

12import numpy as np 

13 

14import sympy 

15import sympy.abc 

16import sympy.physics.quantum 

17from sympy.core.function import AppliedUndef 

18from sympy.core.sympify import converter 

19from sympy.core.core import all_classes as sympy_classes 

20from sympy.physics.matrices import msigma as _msigma 

21 

22import warnings 

23 

24from .._common import reraise_warnings 

25 

26# TODO: remove when sympy correctly includes MutableDenseMatrix (lol). 

27sympy_classes = set(sympy_classes) | {sympy.MutableDenseMatrix} 

28 

29momentum_operators = sympy.symbols('k_x k_y k_z', commutative=False) 

30position_operators = sympy.symbols('x y z', commutative=False) 

31 

32pauli = [sympy.eye(2), _msigma(1), _msigma(2), _msigma(3)] 

33 

34extra_ns = sympy.abc._clash.copy() 

35extra_ns.update({s.name: s for s in momentum_operators}) 

36extra_ns.update({s.name: s for s in position_operators}) 

37extra_ns.update({'kron': sympy.physics.quantum.TensorProduct, 

38 'eye': sympy.eye, 'identity': sympy.eye}) 

39extra_ns.update({'sigma_{}'.format(c): p for c, p in zip('0xyz', pauli)}) 

40 

41 

42# TODO: remove this once https://github.com/sympy/sympy/issues/12060 is fixed 

43del extra_ns['I'] 

44del extra_ns['pi'] 

45 

46 

47################ Helpers to handle sympy 

48 

49def lambdify(expr, locals=None): 

50 """Return a callable object for computing a continuum Hamiltonian. 

51 

52 .. warning:: 

53 This function uses ``eval`` (because it calls ``sympy.sympify``), and 

54 thus should not be used on unsanitized input. 

55 

56 If necessary, the given expression is sympified using 

57 `kwant.continuum.sympify`. It is then converted into a callable object. 

58 

59 Parameters 

60 ---------- 

61 expr : str or SymPy expression 

62 Expression to be converted into a callable object 

63 locals : dict or ``None`` (default) 

64 Additional definitions for `~kwant.continuum.sympify`. 

65 

66 Examples 

67 -------- 

68 >>> f = lambdify('a + b', locals={'b': 'b + c'}) 

69 >>> f(1, 3, 5) 

70 9 

71 

72 >>> ns = {'sigma_plus': [[0, 2], [0, 0]]} 

73 >>> f = lambdify('k_x**2 * sigma_plus', ns) 

74 >>> f(0.25) 

75 array([[ 0. , 0.125], 

76 [ 0. , 0. ]]) 

77 """ 

78 with reraise_warnings(level=4): 

79 expr = sympify(expr, locals) 

80 

81 args = [s.name for s in expr.atoms(sympy.Symbol)] 

82 args += [str(f.func) for f in expr.atoms(AppliedUndef, sympy.Function)] 

83 

84 return sympy.lambdify(sorted(args), expr) 

85 

86 

87def sympify(expr, locals=None): 

88 """Sympify object using special rules for Hamiltonians. 

89 

90 If `'expr`` is already a type that SymPy understands, it will do nothing 

91 but return that value. Note that ``locals`` will not be used in this 

92 situation. 

93 

94 Otherwise, it is sympified by ``sympy.sympify`` with a modified namespace 

95 such that 

96 

97 * the position operators "x", "y" or "z" and momentum operators "k_x", 

98 "k_y", and "k_z" do not commute, 

99 * all single-letter identifiers and names of Greek letters (e.g. "pi" or 

100 "gamma") are treated as symbols, 

101 * "kron" corresponds to ``sympy.physics.quantum.TensorProduct``, and 

102 "identity" to ``sympy.eye``, 

103 * "sigma_0", "sigma_x", "sigma_y", "sigma_z" are the Pauli matrices. 

104 

105 In addition, Python list literals are interpreted as SymPy matrices. 

106 

107 .. warning:: 

108 This function uses ``eval`` (because it calls ``sympy.sympify``), and 

109 thus should not be used on unsanitized input. 

110 

111 Parameters 

112 ---------- 

113 expr : str or SymPy expression 

114 Expression to be converted to a SymPy object. 

115 locals : dict or ``None`` (default) 

116 Additional entries for the namespace under which `expr` is sympified. 

117 The keys must be valid Python variable names. The values may be 

118 strings, since they are all are sent through `continuum.sympify` 

119 themselves before use. (Note that this is a difference to how 

120 ``sympy.sympify`` behaves.) 

121 

122 .. note:: 

123 When a value of `locals` is already a SymPy object, it is used 

124 as-is, and the caller is responsible to set the commutativity of 

125 its symbols appropriately. This possible source of errors is 

126 demonstrated in the last example below. 

127 

128 Returns 

129 ------- 

130 result : SymPy object 

131 

132 Examples 

133 -------- 

134 >>> sympify('k_x * A(x) * k_x + V(x)') 

135 k_x*A(x)*k_x + V(x) # as valid sympy object 

136 

137 >>> sympify('k_x**2 + V', locals={'V': 'V_0 + V(x)'}) 

138 k_x**2 + V(x) + V_0 

139 

140 >>> ns = {'sigma_plus': [[0, 2], [0, 0]]} 

141 >>> sympify('k_x**2 * sigma_plus', ns) 

142 Matrix([ 

143 [0, 2*k_x**2], 

144 [0, 0]]) 

145 

146 >>> sympify('k_x * A(c) * k_x', locals={'c': 'x'}) 

147 k_x*A(x)*k_x 

148 >>> sympify('k_x * A(c) * k_x', locals={'c': sympy.Symbol('x')}) 

149 A(x)*k_x**2 

150 

151 """ 

152 stored_value = None 

153 

154 # if ``expr`` is already a ``sympy`` object we may terminate a code path 

155 if isinstance(expr, tuple(sympy_classes)): 

156 if locals: 

157 warnings.warn('Input expression is already SymPy object: ' 

158 '"locals" will not be used.', 

159 RuntimeWarning, 

160 stacklevel=2) 

161 

162 # We assume that all present functions, like "sin", "cos", will be 

163 # provided by user during the final evaluation through "params". 

164 # Therefore we make sure they are defined as AppliedUndef, not built-in 

165 # sympy types. 

166 subs = {r: sympy.Function(str(r.func))(*r.args) 

167 for r in expr.atoms(sympy.Function)} 

168 

169 return expr.subs(subs) 

170 

171 # if ``expr`` is not a "sympy" then we proceed with sympifying process 

172 if locals is None: 

173 locals = {} 

174 

175 for k in locals: 

176 if (not isinstance(k, str) 

177 or not k.isidentifier() or keyword.iskeyword(k)): 

178 raise ValueError( 

179 "Invalid key in 'locals': {}\nKeys must be " 

180 "identifiers and may not be keywords".format(repr(k))) 

181 

182 # sympify values of locals before updating it with extra_ns 

183 # Cast numpy array values in locals to sympy matrices to make sure they have 

184 # correct format 

185 locals = {k: (sympy.Matrix(v) if isinstance(v, np.ndarray) else sympify(v)) 

186 for k, v in locals.items()} 

187 

188 for k, v in extra_ns.items(): 

189 locals.setdefault(k, v) 

190 try: 

191 stored_value = converter.pop(list, None) 

192 converter[list] = lambda x: sympy.Matrix(x) 

193 hamiltonian = sympy.sympify(expr, locals=locals) 

194 # if input is for example ``[[k_x * A(x) * k_x]]`` after the first 

195 # sympify we are getting list of sympy objects, so we call sympify 

196 # second time to obtain ``sympy`` matrices. 

197 hamiltonian = sympy.sympify(hamiltonian) 

198 finally: 

199 if stored_value is not None: 199 ↛ 200line 199 didn't jump to line 200, because the condition on line 199 was never true

200 converter[list] = stored_value 

201 else: 

202 del converter[list] 

203 

204 return hamiltonian 

205 

206 

207def make_commutative(expr, *symbols): 

208 """Make sure that specified symbols are defined as commutative. 

209 

210 Parameters 

211 ---------- 

212 expr: sympy.Expr or sympy.Matrix 

213 symbols: sequace of symbols 

214 Set of symbols that are requiered to be commutative. It doesn't matter 

215 of symbol is provided as commutative or not. 

216 

217 Returns 

218 ------- 

219 input expression with all specified symbols changed to commutative. 

220 """ 

221 symbols = [sympy.Symbol(s.name, commutative=False) for s in symbols] 

222 expr = expr.subs({s: sympy.Symbol(s.name) for s in symbols}) 

223 return expr 

224 

225 

226def monomials(expr, gens=None): 

227 """Parse ``expr`` into monomials in the symbols in ``gens``. 

228 

229 Parameters 

230 ---------- 

231 expr: sympy.Expr or sympy.Matrix 

232 Sympy expression to be parsed into monomials. 

233 gens: sequence of sympy.Symbol objects or strings (optional) 

234 Generators of monomials. If unset it will default to all 

235 symbols used in ``expr``. 

236 

237 Returns 

238 ------- 

239 dictionary (generator: monomial) 

240 

241 Examples 

242 -------- 

243 >>> expr = kwant.continuum.sympify("A * (x**2 + y) + B * x + C") 

244 >>> monomials(expr, gens=('x', 'y')) 

245 {1: C, x: B, x**2: A, y: A} 

246 """ 

247 if gens is None: 

248 gens = expr.atoms(sympy.Symbol) 

249 else: 

250 gens = [sympify(g) for g in gens] 

251 

252 if not isinstance(expr, sympy.MatrixBase): 

253 return _expression_monomials(expr, gens) 

254 else: 

255 output = defaultdict(lambda: sympy.zeros(*expr.shape)) 

256 for (i, j), e in np.ndenumerate(expr): 

257 mons = _expression_monomials(e, gens) 

258 for key, val in mons.items(): 

259 output[key][i, j] += val 

260 return dict(output) 

261 

262 

263def _expression_monomials(expr, gens): 

264 """Parse ``expr`` into monomials in the symbols in ``gens``. 

265 

266 Parameters 

267 ---------- 

268 expr: sympy.Expr 

269 Sympy expr to be parsed. 

270 gens: sequence of sympy.Symbol 

271 Generators of monomials. 

272 

273 Returns 

274 ------- 

275 dictionary (generator: monomial) 

276 """ 

277 expr = sympy.expand(expr) 

278 output = defaultdict(lambda: sympy.Integer(0)) 

279 for summand in expr.as_ordered_terms(): 

280 key = [] 

281 val = [] 

282 for factor in summand.as_ordered_factors(): 

283 symbol, exponent = factor.as_base_exp() 

284 if symbol in gens: 

285 key.append(factor) 

286 else: 

287 val.append(factor) 

288 output[sympy.Mul(*key)] += sympy.Mul(*val) 

289 

290 return dict(output) 

291 

292 

293################ general help functions 

294 

295def gcd(*args): 

296 if len(args) == 1: 

297 return args[0] 

298 

299 L = list(args) 

300 

301 while len(L) > 1: 

302 a = L[len(L) - 2] 

303 b = L[len(L) - 1] 

304 L = L[:len(L) - 2] 

305 

306 while a: 

307 a, b = b%a, a 

308 

309 L.append(b) 

310 

311 return abs(b)