Coverage for tbkit/neighbours.py: 100%
108 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-15 06:13 +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 :param lat: **Lattice** instance. Only *unit_cell* and *prim_vec* are used.
82 :param n_max: Positive integer. Number of shells.
84 :returns:
85 * **shells** -- Real ndarray, shape (n_max,). Distances, ascending
86 (rounded to 4 decimals, as in *System*).
87 '''
88 error_handling.lat(lat)
89 error_handling.positive_int(n_max, 'n_max')
90 return _shells(lat, n_max)[0]
93def _shells(lat: Lattice, n_max: int, r_max: float | None = None) -> tuple[NDArray, int]:
94 r'''
95 Private function. The first *n_max* shell distances, and the cell range
96 *m* whose candidates hold every bond of those shells: a bond of length
97 :math:`|\mathbf{d}|` has :math:`|n_k| \le |\mathbf{b}_k|(|\mathbf{d}| + \max|\boldsymbol\tau_j-\boldsymbol\tau_i|)/2\pi`,
98 so a box of half-width *m* is complete up to the radius
99 :math:`2\pi m/\max_k|\mathbf{b}_k| - \max|\boldsymbol\tau_j-\boldsymbol\tau_i|`.
100 With *r_max*, every shell up to (at least) that radius instead.
101 '''
102 a, tau, _ = _geometry(lat)
103 b_norm = np.linalg.norm(np.linalg.pinv(a), axis=0).max() # max |b_k| / 2 pi
104 dtau = np.linalg.norm(tau[:, None, :] - tau[None, :, :], axis=2).max()
105 m = 1
106 while True:
107 _, _, _, d = _candidates(lat, m)
108 dist = np.linalg.norm(d, axis=1)
109 radius = m / b_norm - dtau - 2 * ATOL
110 uni = np.unique(dist[dist > ATOL].round(4))
111 uni = uni[uni <= radius]
112 if r_max is not None:
113 if radius > r_max:
114 return uni, m
115 elif len(uni) >= n_max:
116 return uni[:n_max], m
117 m += 1
120def shell_index(lat: Lattice, dist: NDArray[np.float64]) -> NDArray[np.int64]:
121 '''
122 Private function. Neighbour order of bond lengths *dist* (0 for a
123 length below ATOL, as *System*'s ``dist_uni[0]``).
124 '''
125 dist = np.asarray(dist, dtype='f8')
126 if not np.any(dist > ATOL):
127 return np.zeros(len(dist), dtype=int)
128 shells, _ = _shells(lat, 0, r_max=float(dist.max()) + ATOL)
129 n = np.argmin(np.abs(dist[:, None] - shells[None, :]), axis=1) + 1
130 return np.where(dist > ATOL, n, 0)
133def neighbour_bonds(lat: Lattice, n_max: int) -> NDArray:
134 r'''
135 Get every bond of the first *n_max* neighbour shells of the infinite
136 periodic lattice, once each, oriented so that its angle lies in
137 :math:`[0, 180)` (see the module docstring).
139 :param lat: **Lattice** instance. Only *unit_cell* and *prim_vec* are used.
140 :param n_max: Positive integer. Number of shells.
142 :returns:
143 * **bonds** -- Structured ndarray with fields 'n' (shell), 'i', 'j'
144 (orbitals), 'R' (lattice vector, *dim* integers), 'ang' (degrees,
145 in [0, 180)), 'tag' (``tag_i + tag_j``) and 'dis' (length),
146 sorted by (n, i, j, R).
148 Example usage::
150 bonds = neighbour_bonds(lattices.honeycomb(), 1)
151 # three 'ab'/'ba' bonds of length 1, at angles 30, 90 and 150
152 '''
153 error_handling.lat(lat)
154 error_handling.positive_int(n_max, 'n_max')
155 return _bonds(lat, n_max)
158def _bonds(lat: Lattice, n_max: int) -> NDArray:
159 '''
160 Private function. *neighbour_bonds* without the argument checks.
161 '''
162 shells, m = _shells(lat, n_max)
163 i, j, R, d = _candidates(lat, m)
164 dist = np.linalg.norm(d, axis=1)
165 diff = np.abs(dist[:, None] - shells[None, :])
166 n = np.argmin(diff, axis=1)
167 ang = _bond_angles(d)
168 keep = (diff[np.arange(len(dist)), n] <= ATOL) & (dist > ATOL) & _upper(ang)
169 _, _, tags = _geometry(lat)
170 dim = len(lat.prim_vec)
171 bonds = np.zeros(int(keep.sum()), dtype=[('n', 'u2'), ('i', 'u4'), ('j', 'u4'), ('R', 'i8', (dim,)),
172 ('ang', 'f8'), ('tag', 'U2'), ('dis', 'f8')])
173 bonds['n'] = n[keep] + 1
174 bonds['i'], bonds['j'], bonds['R'] = i[keep], j[keep], R[keep]
175 bonds['ang'] = np.maximum(ang[keep], 0.)
176 bonds['tag'] = np.char.add(tags[i[keep]], tags[j[keep]])
177 bonds['dis'] = dist[keep]
178 order = np.lexsort(tuple(bonds['R'][:, k] for k in range(dim)[::-1]) + (bonds['j'], bonds['i'], bonds['n']))
179 return bonds[order]
182def neighbour_hoppings(
183 lat: Lattice, list_hop: list[dict], spin: bool = False, hermitian: bool = True,
184) -> list[dict]:
185 r'''
186 Turn hoppings given by neighbour order, as in *System.set_hopping*,
187 into the explicit ``{'i', 'j', 'R', 't'}`` list of *KSpace.set_hopping*.
189 :param lat: **Lattice** instance. Only *unit_cell* and *prim_vec* are used.
190 :param list_hop: List of dictionaries with keys ('n', 't') and
191 optionally 'ang' and/or 'tag' (see *System.set_hopping*):
193 * 'n': Positive integer. Neighbour order (1 for nearest neighbours).
194 * 'ang': Real number, in degrees. Only the bonds along this angle.
195 An angle in :math:`[0, 180)` selects the bonds oriented as they
196 are stored (see *neighbour_bonds*), an angle in :math:`[-180, 0)`
197 the same bonds oriented the other way (from *j* to *i*).
198 * 'tag': String of length 2. Only the bonds from a site of the
199 first sublattice to one of the second, in the orientation above.
200 * 't': Complex number (or, if *spin*, a 2x2 matrix): the matrix
201 element :math:`\langle i|H|j\rangle` along the oriented bond. Or a
202 value function (see *KSpace.set_hopping*), passed on as it is (as
203 ``t(j, i)^\dagger`` for a reversed bond of a Hermitian model).
205 When several dictionaries select the same bond, the last one wins,
206 as in *System.set_hopping*.
207 :param spin: Boolean. Default value False. See *KSpace*.
208 :param hermitian: Boolean. Default value True. If True, each bond is
209 returned once, in its stored orientation (a hopping given along a
210 negative angle becomes its conjugate transpose); if False, the two
211 orientations are independent hoppings (non-reciprocal models).
213 :returns:
214 * **list_hop** -- List of dictionaries ('i', 'j', 'R', 't').
216 Example usage::
218 # graphene: nearest (t) and next-nearest (t2) neighbours
219 explicit = neighbour_hoppings(lattices.honeycomb(),
220 [{'n': 1, 't': -1.}, {'n': 2, 't': 0.1}])
221 '''
222 error_handling.lat(lat)
223 error_handling.set_hopping_neighbours(list_hop, spin)
224 error_handling.boolean(spin, 'spin')
225 error_handling.boolean(hermitian, 'hermitian')
226 bonds = _bonds(lat, max(dic['n'] for dic in list_hop))
227 chosen = {}
228 for dic in list_hop:
229 shell = bonds[bonds['n'] == dic['n']]
230 reverse = 'ang' in dic and dic['ang'] < 0
231 ind = np.ones(len(shell), bool)
232 if 'ang' in dic:
233 error_handling.shell_angle(dic['ang'], shell['ang'])
234 ind &= np.isclose(dic['ang'] + 180. * reverse, shell['ang'], atol=ATOL)
235 if 'tag' in dic:
236 ind &= shell['tag'] == (dic['tag'][::-1] if reverse else dic['tag'])
237 error_handling.index(ind, dic)
238 t = dic['t']
239 if reverse and hermitian:
240 t_rev = values.reversed_conj(t) if callable(t) else np.conj(t) if np.ndim(t) == 0 else np.asarray(t).conj().T
241 for bond in shell[ind]:
242 i, j, R = int(bond['i']), int(bond['j']), tuple(int(c) for c in bond['R'])
243 if not reverse:
244 chosen[(i, j, R)] = t
245 elif hermitian:
246 chosen[(i, j, R)] = t_rev
247 else:
248 chosen[(j, i, tuple(-c for c in R))] = t
249 return [{'i': i, 'j': j, 'R': R, 't': t} for (i, j, R), t in chosen.items()]