Coverage for tbkit/neighbours.py: 100%
108 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-30 11:04 +0000
1r'''
2Neighbour shells of a periodic lattice, and the neighbour-order hoppings of
3*KSpace.set_hopping*.
5*System.set_hopping* addresses the bonds of a finite lattice by neighbour
6order ``'n'`` (1st, 2nd, ... shortest distance), optionally narrowed by
7bond angle ``'ang'`` and sublattice-pair tag ``'tag'``. This module finds
8the same bonds in a *periodic* lattice, from *unit_cell* and *prim_vec*
9only: each bond is an intra-cell orbital pair :math:`(i, j)` plus a lattice
10vector :math:`\mathbf{R} = n_1\mathbf{a}_1+\dots`, with bond vector
11:math:`\mathbf{d} = \mathbf{R} + \boldsymbol\tau_j - \boldsymbol\tau_i`,
12and turns a neighbour-order hopping list into the explicit
13``{'i', 'j', 'R', 't'}`` list of *KSpace.set_hopping*.
15The conventions are those of *System*: distances are grouped after rounding
16to 4 decimals and matched within ``ATOL = 1e-3``; every bond is kept once,
17oriented from *i* to *j* so that its angle (that of the projection of
18:math:`\mathbf{d}` on the :math:`(x, y)` plane, in degrees; a vertical bond
19has angle 0 pointing up, -180 pointing down) lies in :math:`[0, 180)`; its
20tag is ``tag_i + tag_j``.
21'''
22from __future__ import annotations
24import numpy as np
25from numpy.typing import NDArray
27import tbkit.error_handling as error_handling
28import tbkit.values as values
29from tbkit.lattice import Lattice
30from tbkit.system import ATOL, _upper
33PI = np.pi
36def _geometry(lat: Lattice) -> tuple[NDArray, NDArray, NDArray]:
37 '''
38 Private function. Primitive vectors (dim, space_dim), orbital positions
39 (n_sites, space_dim) and tags of *lat*.
40 '''
41 a = np.array(lat.prim_vec, dtype='f8')
42 tau = np.array([dic['r0'] for dic in lat.unit_cell], dtype='f8')
43 tags = np.array([dic['tag'] for dic in lat.unit_cell])
44 return a, tau, tags
47def _bond_angles(d: NDArray[np.float64]) -> NDArray[np.float64]:
48 '''
49 Private function. Angles (degrees) of the bond vectors *d* (shape
50 (nb, space_dim)), with the vertical-bond rule of *System.get_distances*.
51 '''
52 ang = 180. / PI * np.arctan2(d[:, 1], d[:, 0])
53 if d.shape[1] == 3:
54 vertical = np.hypot(d[:, 0], d[:, 1]) < 1e-6 * np.maximum(np.linalg.norm(d, axis=1), 1.)
55 ang[vertical & (d[:, 2] > 0)] = 0.
56 ang[vertical & (d[:, 2] <= 0)] = -180.
57 return ang
60def _candidates(lat: Lattice, m: int) -> tuple[NDArray, NDArray, NDArray, NDArray]:
61 '''
62 Private function. Every (i, j, R) with :math:`|n_k|\\le m`, and its bond
63 vector: orbitals i, j (nb,), cells R (nb, dim), bonds d (nb, space_dim).
64 '''
65 a, tau, _ = _geometry(lat)
66 dim, ns = len(a), len(tau)
67 cells = np.array(np.meshgrid(*[np.arange(-m, m + 1)] * dim, indexing='ij')).reshape(dim, -1).T
68 i, j = [m.ravel() for m in np.meshgrid(np.arange(ns), np.arange(ns), indexing='ij')]
69 ii = np.repeat(i[None, :], len(cells), axis=0).ravel()
70 jj = np.repeat(j[None, :], len(cells), axis=0).ravel()
71 R = np.repeat(cells, len(i), axis=0)
72 d = R @ a + tau[jj] - tau[ii]
73 return ii, jj, R, d
76def neighbour_shells(lat: Lattice, n_max: int) -> NDArray[np.float64]:
77 r'''
78 Get the distances of the first *n_max* neighbour shells of the infinite
79 periodic lattice (the ``dist_uni[1:]`` of *System*, without its edges).
81 Parameters
82 ----------
83 lat : **Lattice** instance
84 Only *unit_cell* and *prim_vec* are used.
85 n_max : Positive integer
86 Number of shells.
88 Returns
89 -------
90 shells : Real ndarray, shape (n_max,)
91 Distances, ascending
92 (rounded to 4 decimals, as in *System*).
93 '''
94 error_handling.lat(lat)
95 error_handling.positive_int(n_max, 'n_max')
96 return _shells(lat, n_max)[0]
99def _shells(lat: Lattice, n_max: int, r_max: float | None = None) -> tuple[NDArray, int]:
100 r'''
101 Private function. The first *n_max* shell distances, and the cell range
102 *m* whose candidates hold every bond of those shells: a bond of length
103 :math:`|\mathbf{d}|` has :math:`|n_k| \le |\mathbf{b}_k|(|\mathbf{d}| + \max|\boldsymbol\tau_j-\boldsymbol\tau_i|)/2\pi`,
104 so a box of half-width *m* is complete up to the radius
105 :math:`2\pi m/\max_k|\mathbf{b}_k| - \max|\boldsymbol\tau_j-\boldsymbol\tau_i|`.
106 With *r_max*, every shell up to (at least) that radius instead.
107 '''
108 a, tau, _ = _geometry(lat)
109 b_norm = np.linalg.norm(np.linalg.pinv(a), axis=0).max() # max |b_k| / 2 pi
110 dtau = np.linalg.norm(tau[:, None, :] - tau[None, :, :], axis=2).max()
111 m = 1
112 while True:
113 _, _, _, d = _candidates(lat, m)
114 dist = np.linalg.norm(d, axis=1)
115 radius = m / b_norm - dtau - 2 * ATOL
116 uni = np.unique(dist[dist > ATOL].round(4))
117 uni = uni[uni <= radius]
118 if r_max is not None:
119 if radius > r_max:
120 return uni, m
121 elif len(uni) >= n_max:
122 return uni[:n_max], m
123 m += 1
126def shell_index(lat: Lattice, dist: NDArray[np.float64]) -> NDArray[np.int64]:
127 '''
128 Private function. Neighbour order of bond lengths *dist* (0 for a
129 length below ATOL, as *System*'s ``dist_uni[0]``).
130 '''
131 dist = np.asarray(dist, dtype='f8')
132 if not np.any(dist > ATOL):
133 return np.zeros(len(dist), dtype=int)
134 shells, _ = _shells(lat, 0, r_max=float(dist.max()) + ATOL)
135 n = np.argmin(np.abs(dist[:, None] - shells[None, :]), axis=1) + 1
136 return np.where(dist > ATOL, n, 0)
139def neighbour_bonds(lat: Lattice, n_max: int) -> NDArray:
140 r'''
141 Get every bond of the first *n_max* neighbour shells of the infinite
142 periodic lattice, once each, oriented so that its angle lies in
143 :math:`[0, 180)` (see the module docstring).
145 Parameters
146 ----------
147 lat : **Lattice** instance
148 Only *unit_cell* and *prim_vec* are used.
149 n_max : Positive integer
150 Number of shells.
152 Returns
153 -------
154 bonds
155 Structured ndarray with fields 'n' (shell), 'i', 'j'
156 (orbitals), 'R' (lattice vector, *dim* integers), 'ang' (degrees,
157 in [0, 180)), 'tag' (``tag_i + tag_j``) and 'dis' (length),
158 sorted by (n, i, j, R).
160 Examples
161 --------
162 Example usage::
164 bonds = neighbour_bonds(lattices.honeycomb(), 1)
165 # three 'ab'/'ba' bonds of length 1, at angles 30, 90 and 150
166 '''
167 error_handling.lat(lat)
168 error_handling.positive_int(n_max, 'n_max')
169 return _bonds(lat, n_max)
172def _bonds(lat: Lattice, n_max: int) -> NDArray:
173 '''
174 Private function. *neighbour_bonds* without the argument checks.
175 '''
176 shells, m = _shells(lat, n_max)
177 i, j, R, d = _candidates(lat, m)
178 dist = np.linalg.norm(d, axis=1)
179 diff = np.abs(dist[:, None] - shells[None, :])
180 n = np.argmin(diff, axis=1)
181 ang = _bond_angles(d)
182 keep = (diff[np.arange(len(dist)), n] <= ATOL) & (dist > ATOL) & _upper(ang)
183 _, _, tags = _geometry(lat)
184 dim = len(lat.prim_vec)
185 bonds = np.zeros(int(keep.sum()), dtype=[('n', 'u2'), ('i', 'u4'), ('j', 'u4'), ('R', 'i8', (dim,)),
186 ('ang', 'f8'), ('tag', 'U2'), ('dis', 'f8')])
187 bonds['n'] = n[keep] + 1
188 bonds['i'], bonds['j'], bonds['R'] = i[keep], j[keep], R[keep]
189 bonds['ang'] = np.maximum(ang[keep], 0.)
190 bonds['tag'] = np.char.add(tags[i[keep]], tags[j[keep]])
191 bonds['dis'] = dist[keep]
192 order = np.lexsort(tuple(bonds['R'][:, k] for k in range(dim)[::-1]) + (bonds['j'], bonds['i'], bonds['n']))
193 return bonds[order]
196def neighbour_hoppings(
197 lat: Lattice, list_hop: list[dict], spin: bool = False, hermitian: bool = True,
198) -> list[dict]:
199 r'''
200 Turn hoppings given by neighbour order, as in *System.set_hopping*,
201 into the explicit ``{'i', 'j', 'R', 't'}`` list of *KSpace.set_hopping*.
203 Parameters
204 ----------
205 lat : **Lattice** instance
206 Only *unit_cell* and *prim_vec* are used.
207 list_hop
208 List of dictionaries with keys ('n', 't') and
209 optionally 'ang' and/or 'tag' (see *System.set_hopping*):
211 * 'n': Positive integer. Neighbour order (1 for nearest neighbours).
212 * 'ang': Real number, in degrees. Only the bonds along this angle.
213 An angle in :math:`[0, 180)` selects the bonds oriented as they
214 are stored (see *neighbour_bonds*), an angle in :math:`[-180, 0)`
215 the same bonds oriented the other way (from *j* to *i*).
216 * 'tag': String of length 2. Only the bonds from a site of the
217 first sublattice to one of the second, in the orientation above.
218 * 't': Complex number (or, if *spin*, a 2x2 matrix): the matrix
219 element :math:`\langle i|H|j\rangle` along the oriented bond. Or a
220 value function (see *KSpace.set_hopping*), passed on as it is (as
221 ``t(j, i)^\dagger`` for a reversed bond of a Hermitian model).
223 When several dictionaries select the same bond, the last one wins,
224 as in *System.set_hopping*.
225 spin : Boolean, default False
226 See *KSpace*.
227 hermitian : Boolean, default True
228 If True, each bond is
229 returned once, in its stored orientation (a hopping given along a
230 negative angle becomes its conjugate transpose); if False, the two
231 orientations are independent hoppings (non-reciprocal models).
233 Returns
234 -------
235 list_hop : List of dictionaries ('i', 'j', 'R', 't')
237 Examples
238 --------
239 Example usage::
241 # graphene: nearest (t) and next-nearest (t2) neighbours
242 explicit = neighbour_hoppings(lattices.honeycomb(),
243 [{'n': 1, 't': -1.}, {'n': 2, 't': 0.1}])
244 '''
245 error_handling.lat(lat)
246 error_handling.set_hopping_neighbours(list_hop, spin)
247 error_handling.boolean(spin, 'spin')
248 error_handling.boolean(hermitian, 'hermitian')
249 bonds = _bonds(lat, max(dic['n'] for dic in list_hop))
250 chosen = {}
251 for dic in list_hop:
252 shell = bonds[bonds['n'] == dic['n']]
253 reverse = 'ang' in dic and dic['ang'] < 0
254 ind = np.ones(len(shell), bool)
255 if 'ang' in dic:
256 error_handling.shell_angle(dic['ang'], shell['ang'])
257 ind &= np.isclose(dic['ang'] + 180. * reverse, shell['ang'], atol=ATOL)
258 if 'tag' in dic:
259 ind &= shell['tag'] == (dic['tag'][::-1] if reverse else dic['tag'])
260 error_handling.index(ind, dic)
261 t = dic['t']
262 if reverse and hermitian:
263 t_rev = values.reversed_conj(t) if callable(t) else np.conj(t) if np.ndim(t) == 0 else np.asarray(t).conj().T
264 for bond in shell[ind]:
265 i, j, R = int(bond['i']), int(bond['j']), tuple(int(c) for c in bond['R'])
266 if not reverse:
267 chosen[(i, j, R)] = t
268 elif hermitian:
269 chosen[(i, j, R)] = t_rev
270 else:
271 chosen[(j, i, tuple(-c for c in R))] = t
272 return [{'i': i, 'j': j, 'R': R, 't': t} for (i, j, R), t in chosen.items()]