Coverage for tbkit/continuum.py: 100%
111 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-15 06:13 +0000
1r'''
2Continuum discretization: turn a k·p (continuum) Hamiltonian, a polynomial
3in the momenta :math:`k_x, k_y, k_z` with constant (possibly symbolic)
4matrix coefficients, into a Tight-Binding model on a square or cubic grid
5of spacing *a*.
7Each momentum is the derivative :math:`k_x = -i\partial_x`, and a power
8:math:`k_x^n` is replaced by the half-step finite difference
9:math:`D_hf(x) = [f(x+h/2) - f(x-h/2)]/h` applied *n* times, with
10:math:`h = a` for an even *n* and :math:`h = 2a` for an odd one, so that
11every shift is a whole number of sites (the scheme of Kwant's
12``kwant.continuum``). In k-space this replaces
14.. math::
16 k^n \to \Big(\tfrac{2}{a}\sin\tfrac{ka}{2}\Big)^n \;(n\text{ even}),\qquad
17 k^n \to \Big(\tfrac{1}{a}\sin ka\Big)^n \;(n\text{ odd}),
19which agrees with :math:`k^n` up to :math:`O(a^2)` and keeps a Hermitian
20Hamiltonian Hermitian: :math:`k_x\to-i(\psi_{n+1}-\psi_{n-1})/2a`,
21:math:`k_x^2\to-(\psi_{n+1}-2\psi_n+\psi_{n-1})/a^2`. Products of momenta
22along different axes multiply their stencils. The coefficients do not
23depend on the position, so the order of the factors does not matter.
25* *discretize_symbolic* returns the hopping matrices :math:`T(\mathbf{R})`
26 as sympy matrices in the grid spacing *a*, to inspect.
27* *discretize* builds the **KSpace** model with
28 :math:`H(\mathbf{k}) = \sum_\mathbf{R} T(\mathbf{R})e^{i\mathbf{k}\cdot\mathbf{R}}`.
29 Its free symbols (*M*, *B*, ...) become parameters of value functions,
30 set by ``get_ham(k, M=...)`` or *KSpace.set_params*.
31 *bridges.finite_system* then gives the matching real-space **System**.
33sympy is an optional dependency: ``pip install tbkit[continuum]``.
34'''
35from __future__ import annotations
37import inspect
38import itertools
39import string
40from math import comb
42import numpy as np
44import tbkit.error_handling as error_handling
45from tbkit.kspace import KSpace
46from tbkit.lattice import Lattice
48#: Names of the momenta, along x, y and z.
49MOMENTA = ('k_x', 'k_y', 'k_z')
52def _sympy():
53 '''
54 Private function. The sympy module, imported on first use.
55 '''
56 try:
57 import sympy
58 except ImportError:
59 sympy = None
60 error_handling.sympy_module(sympy)
61 return sympy
64def _namespace(sympy) -> dict:
65 '''
66 Private function. Names available in a string Hamiltonian: the Pauli
67 matrices *sigma_0*, *sigma_x*, *sigma_y*, *sigma_z*, *kron* (Kronecker
68 product), and plain symbols for the names sympy would otherwise read
69 as its own objects (E, N, O, Q, S, beta, gamma, zeta).
70 '''
71 space = {name: sympy.Symbol(name) for name in ('E', 'N', 'O', 'Q', 'S', 'beta', 'gamma', 'zeta')}
72 space.update({'sigma_0': sympy.eye(2),
73 'sigma_x': sympy.Matrix([[0, 1], [1, 0]]),
74 'sigma_y': sympy.Matrix([[0, -sympy.I], [sympy.I, 0]]),
75 'sigma_z': sympy.Matrix([[1, 0], [0, -1]]),
76 'kron': lambda *ms: sympy.kronecker_product(*ms)})
77 return space
80def _matrix(sympy, hamiltonian):
81 '''
82 Private function. The Hamiltonian as a square sympy Matrix whose free
83 symbols are all real, and the map back to the original symbols.
84 '''
85 error_handling.continuum_type(hamiltonian, (str, sympy.Basic, sympy.MatrixBase))
86 ham = hamiltonian
87 if isinstance(hamiltonian, str):
88 from sympy.parsing.sympy_parser import parse_expr
89 error = None
90 try:
91 ham = parse_expr(hamiltonian, local_dict=_namespace(sympy))
92 except Exception as err: # parse_expr raises many types
93 error = err
94 error_handling.continuum_parsed(error, hamiltonian)
95 error_handling.continuum_type(ham, (sympy.Basic, sympy.MatrixBase))
96 ham = sympy.Matrix(ham) if isinstance(ham, sympy.MatrixBase) else sympy.Matrix([[ham]])
97 error_handling.continuum_square(ham.shape)
98 real = {s: sympy.Symbol(s.name, real=True) for s in ham.free_symbols}
99 return ham.xreplace(real), {r: s for s, r in real.items()}
102def _stencil(sympy, n: int) -> dict[int, object]:
103 r'''
104 Private function. The finite difference of :math:`k^n`, times
105 :math:`a^n`: {shift in sites: coefficient}.
106 '''
107 h = 1 if n % 2 == 0 else 2
108 return {(n - 2*m) * h // 2: (-sympy.I / h)**n * comb(n, m) * (-1)**m for m in range(n + 1)}
111def _discretize(sympy, hamiltonian, dim: int | None):
112 '''
113 Private function. The hoppings of *discretize_symbolic*, the dimension,
114 the number of orbitals, and whether the Hamiltonian is Hermitian (for
115 real parameters).
116 '''
117 ham, originals = _matrix(sympy, hamiltonian)
118 names = {s.name for s in ham.free_symbols}
119 error_handling.continuum_symbols(names)
120 used = [d for d, k in enumerate(MOMENTA) if k in names]
121 if dim is None:
122 dim = max(used, default=0) + 1
123 error_handling.continuum_dim(dim, used)
124 momenta = [sympy.Symbol(k, real=True) for k in MOMENTA[:dim]]
125 a = sympy.Symbol('a')
126 norb = ham.shape[0]
127 hops: dict = {}
128 for i, j in itertools.product(range(norb), repeat=2):
129 entry = sympy.expand(ham[i, j])
130 error_handling.continuum_polynomial((i, j), entry.is_polynomial(*momenta))
131 for powers, coeff in sympy.Poly(entry, *momenta).terms():
132 stencils = [_stencil(sympy, n).items() for n in powers]
133 for steps in itertools.product(*stencils):
134 R = tuple(shift for shift, _ in steps)
135 c = coeff * sympy.Mul(*(c for _, c in steps)) / a**sum(powers)
136 hops.setdefault(R, sympy.zeros(norb))[i, j] += c
137 hops = {R: m.applyfunc(sympy.expand).xreplace(originals) for R, m in sorted(hops.items())}
138 hops = {R: m for R, m in hops.items() if not m.is_zero_matrix}
139 hermitian = (ham - ham.H).applyfunc(sympy.simplify).is_zero_matrix
140 return hops, dim, norb, bool(hermitian)
143def discretize_symbolic(hamiltonian, dim: int | None = None) -> dict:
144 r'''
145 Discretize a continuum Hamiltonian on a square or cubic grid, keeping
146 everything symbolic: the hopping matrices
147 :math:`T(\mathbf{R}) = \langle\mathbf{0}|H|\mathbf{R}\rangle` between
148 a grid site and the one :math:`\mathbf{R}` sites further, in terms of
149 the grid spacing *a* and the parameters of the Hamiltonian (see the
150 module docstring for the scheme).
152 :param hamiltonian: sympy expression or Matrix, or a string, in the
153 momenta ``k_x``, ``k_y``, ``k_z``: a polynomial in them with
154 constant coefficients. Any other symbol is a parameter, taken real
155 (the symbols ``x``, ``y``, ``z`` and ``a`` are not allowed). A
156 string may use the Pauli matrices ``sigma_0``, ``sigma_x``,
157 ``sigma_y``, ``sigma_z`` and ``kron`` (Kronecker product); ``I``
158 is the imaginary unit.
159 :param dim: 1, 2 or 3. Default value None: the last axis whose momentum
160 appears (1 if none).
162 :returns:
163 * **hoppings** -- Dictionary. key: tuple of *dim* integers
164 :math:`\mathbf{R}`, val: sympy Matrix :math:`T(\mathbf{R})`, in
165 the symbol ``a`` and those of *hamiltonian*.
167 Example usage::
169 discretize_symbolic('k_x**2')
170 # {(-1,): Matrix([[-1/a**2]]), (0,): Matrix([[2/a**2]]), (1,): Matrix([[-1/a**2]])}
171 '''
172 return _discretize(_sympy(), hamiltonian, dim)[0]
175def _value(sympy, expr):
176 '''
177 Private function. A value function returning the matrix element *expr*,
178 whose free symbols are its keyword parameters.
179 '''
180 symbols = sorted(expr.free_symbols, key=lambda s: s.name)
181 names = [s.name for s in symbols]
182 func = sympy.lambdify(symbols, expr, modules='numpy')
184 def value(site_i, site_j, **params):
185 return complex(func(*(params[n] for n in names)))
186 kind = inspect.Parameter
187 value.__signature__ = inspect.Signature(
188 [kind('site_i', kind.POSITIONAL_OR_KEYWORD), kind('site_j', kind.POSITIONAL_OR_KEYWORD)]
189 + [kind(n, kind.KEYWORD_ONLY) for n in names])
190 value.__name__ = str(expr)
191 return value
194def discretize(hamiltonian, a: float = 1., dim: int | None = None,
195 tags: str | list[str] | None = None) -> KSpace:
196 r'''
197 Discretize a continuum (k·p) Hamiltonian into a Tight-Binding **KSpace**
198 model on a square (2D) or cubic (3D) grid, or a chain (1D), of spacing
199 *a*. The grid has one site per unit cell carrying the *norb* orbitals of
200 the Hamiltonian's matrix, all at the origin of the cell, so that
201 :math:`H(\mathbf{k}) = \sum_\mathbf{R} T(\mathbf{R})e^{i\mathbf{k}\cdot\mathbf{R}}`
202 with the hoppings of *discretize_symbolic*. It agrees with the
203 continuum Hamiltonian up to :math:`O(a^2)` near
204 :math:`\mathbf{k}=\mathbf{0}`.
206 A matrix element with no free symbol becomes a number. One with
207 parameters becomes a value function (see *KSpace.set_hopping*): pass
208 the parameters to ``get_ham(k, M=...)``, or to *KSpace.set_params* for
209 the other methods. A Hamiltonian that is Hermitian for real parameters
210 is stored with ``hermitian=True`` (one representative per bond);
211 otherwise every matrix element is stored, with ``hermitian=False``.
212 Constant onsite terms are in *ks.onsite*.
214 :param hamiltonian: sympy expression or Matrix, or a string (see
215 *discretize_symbolic*).
216 :param a: Real number, at least :math:`\sqrt{0.1}\approx 0.316` (the
217 shortest primitive vector of a **Lattice**). Default value 1. Grid
218 spacing: to refine the grid, measure lengths in a smaller unit.
219 :param dim: 1, 2 or 3. Default value None: the last axis whose momentum
220 appears (1 if none).
221 :param tags: String or list of one-character strings. Default value
222 None: 'a', 'b', 'c', ... Tag of each orbital.
224 :returns:
225 * **ks** -- **KSpace** instance (spinless, *norb* orbitals). Its
226 lattice has *prim_vec* :math:`a\hat{\mathbf{x}}, a\hat{\mathbf{y}},
227 (a\hat{\mathbf{z}})`.
229 Example usage::
231 # the BHZ model of HgTe quantum wells, one spin block
232 bhz = discretize('A*(k_x*sigma_x - k_y*sigma_y)'
233 '+ (M - B*(k_x**2 + k_y**2))*sigma_z', a=0.5)
234 bhz.set_params(A=1., B=1., M=1.)
235 bhz.chern_number([0])
236 sys = finite_system(bhz, (40, 40)) # a real-space flake
237 '''
238 error_handling.positive_real(a, 'a')
239 error_handling.continuum_spacing(a)
240 sympy = _sympy()
241 hops, dim, norb, hermitian = _discretize(sympy, hamiltonian, dim)
242 error_handling.continuum_tags(tags, norb)
243 tags = list(tags) if tags is not None else list(string.ascii_letters[:norb])
244 space_dim = max(dim, 2)
245 origin = (0.,) * space_dim
246 prim_vec = [tuple(float(a) * (c == d) for c in range(space_dim)) for d in range(dim)]
247 ks = KSpace(Lattice(unit_cell=[{'tag': t, 'r0': origin} for t in tags], prim_vec=prim_vec))
248 a_sym = sympy.Symbol('a')
249 list_hop = []
250 for R, mat in hops.items():
251 mat = mat.subs(a_sym, a)
252 for i, j in itertools.product(range(norb), repeat=2):
253 expr = mat[i, j]
254 if expr == 0:
255 continue
256 nonzero = [n for n in R if n != 0]
257 onsite = not nonzero and i == j
258 if hermitian and not onsite and not ((nonzero and nonzero[0] > 0) or (not nonzero and i < j)):
259 continue # the conjugate of a kept representative
260 t = _value(sympy, expr) if expr.free_symbols else complex(expr)
261 if not onsite:
262 list_hop.append({'i': i, 'j': j, 'R': R, 't': t})
263 elif callable(t):
264 # a parametrized onsite energy: a value function on the
265 # bond i -> i, R = 0, which set_hopping does not take
266 ks._hop_values.append((np.array([i]), np.array([i]), np.zeros((1, ks.space_dim)), t, False))
267 else:
268 ks.onsite[i] = t
269 if list_hop:
270 ks.set_hopping(list_hop, hermitian=hermitian)
271 if not hermitian:
272 ks._nonreciprocal = True
273 return ks