Coverage for tbkit/lattice.py: 100%
227 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-15 06:13 +0000
1from __future__ import annotations
3from typing import Sequence
5import numpy as np
6from numpy.typing import NDArray
7import matplotlib.pyplot as plt
8from matplotlib.figure import Figure
9from scipy.spatial import cKDTree
10import tbkit.error_handling as error_handling
13PI = np.pi
14COOR_DTYPE = [('x', 'f8'), ('y', 'f8'), ('tag', 'U1')]
15#: Site dtype of a lattice in 3D space (*unit_cell* and *prim_vec* given as 3-tuples).
16COOR_DTYPE_3D = [('x', 'f8'), ('y', 'f8'), ('z', 'f8'), ('tag', 'U1')]
19#################################
20# CLASS LATTICE
21#################################
24class Lattice():
25 r'''
26 Build up 1D, 2D or 3D lattice.
27 Lattice is defined by the discrete operation:
29 .. math::
31 \mathbf{R} = n_1\mathbf{a}_1 + n_2\mathbf{a}_2 + n_3\mathbf{a}_3
33 where :math:`\mathbf{a}_1`, :math:`\mathbf{a}_2`, :math:`\mathbf{a}_3` are the
34 primitive vectors and :math:`n_1`, :math:`n_2`, :math:`n_3` are the number of
35 unit cells along them.
37 :param unit_cell: List of dictionaries.
38 One dictionary per site within the unit cell. Each dictionary has two keys:
40 * 'tag', one-character string. Label of the associated sublattice.
41 * 'r0', Tuple. Position: :math:`(x, y)`, or :math:`(x, y, z)` for a
42 lattice in 3D space.
43 :param prim_vec: List of tuples.
44 Define the primitive vectors. List of one/two/three tuples for 1D/2D/3D
45 lattices, each the cartesian coordinates of a primitive vector: 2-tuples
46 in the plane, or 3-tuples in 3D space (then *r0* must be 3-tuples too,
47 and *coor* gains a 'z' field, see :data:`COOR_DTYPE_3D`).
49 In 3D space, the geometry methods acting in the plane (*rotation*,
50 *ellipse_in*, *boundary_line*, ...) act on :math:`(x, y)` only (rotation
51 about the z axis, cylinders, vertical planes), and *plot* draws the
52 projection on the :math:`(x, y)` plane.
54 Example usage::
56 # Line-Centered Square lattice
57 unit_cell = [{'tag': 'a', 'r0': (0., 0.)}, {'tag': 'a', 'r0': (0., 1.)}]
58 prim_vec = [(0, 2), (2, 0)]
59 lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec)
60 '''
62 def __init__(self, unit_cell: list[dict], prim_vec: Sequence[tuple[float, ...]]) -> None:
63 error_handling.unit_cell(unit_cell)
64 error_handling.prim_vec(prim_vec)
65 error_handling.space_dim(unit_cell, prim_vec)
66 self.unit_cell = unit_cell
67 self.prim_vec = prim_vec
68 self.tags = np.unique(np.array([dic['tag'] for dic in self.unit_cell]))
69 self.space_dim = len(prim_vec[0]) # 2, or 3 for a lattice in 3D space
70 self.dtype = COOR_DTYPE if self.space_dim == 2 else COOR_DTYPE_3D
71 self.n1, self.n2, self.n3 = 0, 0, 0
72 self.coor = np.array([], dtype=self.dtype)
73 self.sites = 0
75 def get_lattice(self, n1: int, n2: int = 1, n3: int = 1) -> None:
76 r'''
77 Get the lattice positions, sorted by (y, x) -- by (z, y, x) in 3D space.
79 :param n1: Positive Integer.
80 Number of unit cells along :math:`\mathbf{a}_1`.
81 :param n2: Positive Integer. Default value 1.
82 Number of unit cells along :math:`\mathbf{a}_2`.
83 :param n3: Positive Integer. Default value 1.
84 Number of unit cells along :math:`\mathbf{a}_3` (3D lattices only).
86 Example usage::
88 # Line-Centered Square lattice
89 unit_cell = [{'tag': 'a', 'r0': (0., 0.)}, {'tag': 'a', 'r0': (0., 1.)}]
90 prim_vec = [(0, 2), (2, 0)]
91 lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec)
92 lat.get_lattice(n1=4, n2=5)
93 '''
94 error_handling.get_lattice(self.prim_vec, n1, n2, n3)
95 sites_uc = len(self.unit_cell)
96 sites_tag = n1*n2*n3
97 self.sites = sites_uc * sites_tag
98 self.coor = np.empty(self.sites, dtype=self.dtype)
99 self.n1, self.n2, self.n3 = n1, n2, n3
100 # translations R = i1 a1 + i2 a2 (+ i3 a3), i1 running fastest
101 axes = [f for f in ('x', 'y', 'z')[:self.space_dim]]
102 trans = np.zeros((sites_tag, self.space_dim))
103 counts = (n1, n2, n3)
104 index = np.indices(counts[:max(len(self.prim_vec), 1)][::-1]).reshape(
105 max(len(self.prim_vec), 1), -1)[::-1]
106 for n, a in zip(index, self.prim_vec):
107 trans += n[:, None] * np.array(a, dtype='f8')[None, :]
108 for i, dic in enumerate(self.unit_cell):
109 for c, f in enumerate(axes):
110 self.coor[f][i*sites_tag: (i+1)*sites_tag] = trans[:, c] + dic['r0'][c]
111 self.coor['tag'][i*sites_tag: (i+1)*sites_tag] = dic['tag']
112 self.coor = np.sort(self.coor, order=self.sort_order())
114 def sort_order(self) -> tuple[str, ...]:
115 '''
116 Private method. Fields the sites are sorted by: ('y', 'x'), or
117 ('z', 'y', 'x') in 3D space.
118 '''
119 return ('y', 'x') if self.space_dim == 2 else ('z', 'y', 'x')
121 def distances(self) -> tuple[NDArray, ...]:
122 '''
123 Private method. Pairwise coordinate differences, *d[i, j]* from site
124 *i* to site *j*: (dx, dy), or (dx, dy, dz) in 3D space.
125 '''
126 return tuple(self.coor[f] - self.coor[f].reshape(self.sites, 1)
127 for f in ('x', 'y', 'z')[:self.space_dim])
129 def add_sites(self, coor: NDArray) -> None:
130 '''
131 Add sites.
133 :param coor: Structured array with keys: {'x', 'y', 'tag'}
134 ({'x', 'y', 'z', 'tag'} for a lattice in 3D space).
136 Example usage::
138 # Square lattice
139 unit_cell = [{'tag': 'a', 'r0': (0., 0.)}]
140 prim_vec = [(0, 1), (1, 0)]
141 lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec)
142 lat.get_lattice(n1=2, n2=2)
143 coor = np.array([(-1., -1, 'b'), (-2., -2, 'c')],
144 dtype=[('x', 'f8'), ('y', 'f8'), ('tag', 'U1')])
145 lat.add_sites(coor)
146 '''
147 error_handling.coor(coor, self.dtype)
148 self.coor = np.concatenate([self.coor, coor])
149 self.sites += len(coor)
150 self.tags = np.unique(np.concatenate([self.tags, coor['tag']]))
151 self.coor = np.sort(self.coor, order=self.sort_order())
153 def remove_sites(self, index: list[int]) -> None:
154 '''
155 Remove sites defined by their indices
156 (use method lattice.plot(plt_index=True)
157 to get access to the site indices).
159 :param index: List. Site indices to be removed.
161 Example usage::
163 # Square lattice
164 unit_cell = [{'tag': 'a', 'r0': (0., 0.)}]
165 prim_vec = [(0, 1), (1, 0)]
166 lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec)
167 lat.get_lattice(n1=2, n2=2)
168 lat.remove_sites([0, 2])
169 '''
170 error_handling.empty_coor(self.coor)
171 error_handling.remove_sites(index, self.sites)
172 mask = np.ones(self.sites, bool)
173 mask[index] = False
174 self.coor = self.coor[mask]
175 self.sites = self.coor.size
177 def remove_dangling(self) -> None:
178 '''
179 Remove dangling sites
180 (sites connected with just another site).
181 '''
182 error_handling.empty_coor(self.coor)
183 while True:
184 coords = np.stack([self.coor[f] for f in ('x', 'y', 'z')[:self.space_dim]], axis=1)
185 tree = cKDTree(coords)
186 near = tree.query(coords, k=min(self.sites, 8))[0][:, 1:]
187 len_hop = np.min(near[near > 0]) # the shortest distance between sites
188 pairs = tree.query_pairs(len_hop * (1. + 1e-5) + 1e-8, output_type='ndarray')
189 dis = np.linalg.norm(coords[pairs[:, 1]] - coords[pairs[:, 0]], axis=1)
190 pairs = pairs[np.isclose(dis, len_hop)]
191 degree = np.bincount(pairs.ravel(), minlength=self.sites)
192 dang = list(np.flatnonzero(degree == 1))
193 self.coor = np.delete(self.coor, dang, axis=0)
194 self.sites -= len(dang)
195 if dang == []:
196 break
198 def shift_x(self, shift: float) -> None:
199 '''
200 Shift the x coordinates.
202 :param shift: Real number. Shift value.
203 '''
204 error_handling.empty_coor(self.coor)
205 error_handling.real_number(shift, 'shift')
206 self.coor['x'] += shift
208 def shift_y(self, shift: float) -> None:
209 '''
210 Shift the y coordinates.
212 :param shift: Real number. Shift value.
213 '''
214 error_handling.empty_coor(self.coor)
215 error_handling.real_number(shift, 'shift')
216 self.coor['y'] += shift
218 def change_sign_x(self) -> None:
219 '''
220 Change x coordinates sign.
221 '''
222 error_handling.empty_coor(self.coor)
223 self.coor['x'] *= -1
225 def change_sign_y(self) -> None:
226 '''
227 Change y coordinates sign.
228 '''
229 error_handling.empty_coor(self.coor)
230 self.coor['y'] *= -1
232 def shift_z(self, shift: float) -> None:
233 '''
234 Shift the z coordinates (3D lattices only).
236 :param shift: Real number. Shift value.
237 '''
238 error_handling.empty_coor(self.coor)
239 error_handling.space_3d(self.space_dim)
240 error_handling.real_number(shift, 'shift')
241 self.coor['z'] += shift
243 def change_sign_z(self) -> None:
244 '''
245 Change z coordinates sign (3D lattices only).
246 '''
247 error_handling.empty_coor(self.coor)
248 error_handling.space_3d(self.space_dim)
249 self.coor['z'] *= -1
251 def slab(self, z_min: float, z_max: float) -> None:
252 r'''
253 Keep only the sites with :math:`z_{min} < z < z_{max}` (3D lattices only).
255 :param z_min: Real number.
256 :param z_max: Real number, larger than *z_min*.
257 '''
258 error_handling.empty_coor(self.coor)
259 error_handling.space_3d(self.space_dim)
260 error_handling.real_number(z_min, 'z_min')
261 error_handling.real_number(z_max, 'z_max')
262 error_handling.smaller(z_min, 'z_min', z_max, 'z_max')
263 self.coor = self.coor[(self.coor['z'] > z_min) & (self.coor['z'] < z_max)]
264 self.sites = len(self.coor)
266 def boundary_line(self, cx: float, cy: float, co: float) -> None:
267 r'''
268 Keep only the sites with :math:`c_yy+c_xx > c_0`.
270 :param cx: Real number. :math:`c_x` value.
271 :param cy: Real number. :math:`c_y` value.
272 :param co: Real number. :math:`c_0` value.
273 '''
274 error_handling.empty_coor(self.coor)
275 error_handling.real_number(cx, 'cx')
276 error_handling.real_number(cy, 'cy')
277 error_handling.real_number(co, 'co')
278 self.coor = self.coor[cy * self.coor['y'] + cx * self.coor['x'] > co]
279 self.sites = len(self.coor)
281 def ellipse_in(self, rx: float, ry: float, x0: float, y0: float) -> None:
282 r'''
283 Select sites according to
285 .. math::
287 (x-x_0)^2/r_x^2+(y-y_0)^2/r_y^2 < 1\, .
289 :param rx: Positive Real number. Radius along :math:`x`.
290 :param ry: Positive Real number. Radius along :math:`y`.
291 :param x0: Real number. :math:`x` center.
292 :param y0: Real number. :math:`y` center.
293 '''
294 error_handling.empty_coor(self.coor)
295 error_handling.positive_real(rx, 'rx')
296 error_handling.positive_real(ry, 'ry')
297 error_handling.real_number(x0, 'x0')
298 error_handling.real_number(y0, 'y0')
299 self.coor = self.coor[(self.coor['x'] -x0) ** 2 / rx ** 2 + \
300 (self.coor['y'] -y0) ** 2 / ry ** 2 < 1.]
301 self.sites = len(self.coor)
303 def ellipse_out(self, rx: float, ry: float, x0: float, y0: float) -> None:
304 r'''
305 Select sites according to
307 .. math::
309 (x-x_0)^2/r_x^2+(y-y_0)^2/r_y^2 > 1\, .
312 :param rx: Positive Real number. Radius along :math:`x`.
313 :param ry: Positive Real number. Radius along :math:`y`.
314 :param x0: Real number. :math:`x` center.
315 :param y0: Real number. :math:`y` center.
316 '''
317 error_handling.empty_coor(self.coor)
318 error_handling.positive_real(rx, 'rx')
319 error_handling.positive_real(ry, 'ry')
320 error_handling.real_number(x0, 'x0')
321 error_handling.real_number(y0, 'y0')
322 self.coor = self.coor[(self.coor['x'] -x0) ** 2 / rx ** 2 + \
323 (self.coor['y'] -y0) ** 2 / ry ** 2 > 1.]
324 self.sites = len(self.coor)
326 def center(self) -> None:
327 '''
328 Fix the center of mass of the lattice at (0, 0) -- (0, 0, 0) in 3D space.
329 '''
330 error_handling.empty_coor(self.coor)
331 for f in ('x', 'y', 'z')[:self.space_dim]:
332 self.coor[f] -= np.mean(self.coor[f])
334 def rotation(self, theta: float) -> None:
335 r'''
336 Rotate the lattice structure about the origin (about the z axis, in
337 3D space) by the angle :math:`\theta`.
339 :param theta: Rotation angle in degrees.
340 '''
341 error_handling.empty_coor(self.coor)
342 error_handling.real_number(theta, 'theta')
343 theta *= PI / 180
344 x = self.coor['x'].copy()
345 y = self.coor['y'].copy()
346 self.coor['x'] = x * np.cos(theta) - y * np.sin(theta)
347 self.coor['y'] = y * np.cos(theta) + x * np.sin(theta)
349 def clean_coor(self) -> None:
350 '''
351 Keep only the sites with different coordinates.
352 '''
353 error_handling.empty_coor(self.coor)
354 fields = list(('x', 'y', 'z')[:self.space_dim])
355 coor = self.coor[fields].copy()
356 for f in fields:
357 coor[f] = self.coor[f].round(4)
358 _, idx = np.unique(coor, return_index=True)
359 self.coor = self.coor[idx]
360 self.sites = len(self.coor)
362 def __add__(self, other: 'Lattice') -> 'Lattice':
363 '''
364 Overloading operator +.
365 '''
366 error_handling.lat(other)
367 error_handling.empty_coor(self.coor)
368 error_handling.empty_coor(other.coor)
369 coor = np.concatenate([self.coor, other.coor])
370 tags = np.concatenate([self.tags, other.tags])
371 lat = lattice(unit_cell=self.unit_cell, prim_vec=self.prim_vec)
372 lat.add_sites(coor)
373 lat.sites = self.sites + other.sites
374 lat.tags = np.unique(tags)
375 return lat
377 def __iadd__(self, other: 'Lattice') -> 'Lattice':
378 '''
379 Overloading operator +=.
380 '''
381 error_handling.lat(other)
382 error_handling.empty_coor(self.coor)
383 error_handling.empty_coor(other.coor)
384 self.coor = np.concatenate([self.coor, other.coor])
385 self.sites += other.sites
386 self.tags = np.unique(np.concatenate([self.tags, other.tags]))
387 return self
389 def __sub__(self, other: 'Lattice') -> 'Lattice':
390 '''
391 Overloading operator -.
393 .. note::
395 The tags are not considered in the lattice subtraction.
396 '''
397 error_handling.lat(other)
398 error_handling.empty_coor(self.coor)
399 error_handling.empty_coor(other.coor)
400 boo = self.overlaps(other)
401 coor = self.coor[np.logical_not(boo)]
402 lat = lattice(unit_cell=self.unit_cell, prim_vec=self.prim_vec)
403 lat.add_sites(coor)
404 lat.sites = len(lat.coor)
405 lat.tags = self.tags
406 return lat
408 def __isub__(self, other: 'Lattice') -> 'Lattice':
409 '''
410 Overloading operator -=.
412 .. note::
414 The tags are not considered in the lattice subtraction.
415 '''
416 error_handling.lat(other)
417 error_handling.empty_coor(self.coor)
418 error_handling.empty_coor(other.coor)
419 boo = self.overlaps(other)
420 self.coor = self.coor[np.logical_not(boo)]
421 self.sites = int(np.sum(np.logical_not(boo)))
422 return self
424 def overlaps(self, other: 'Lattice') -> NDArray[np.bool_]:
425 '''
426 Private method. Mask of the sites of *self* that coincide with a site
427 of *other*.
428 '''
429 names = self.coor.dtype.names or ()
430 fields = [f for f in ('x', 'y', 'z') if f in names]
431 boo = np.zeros(self.sites, bool)
432 for c in other.coor:
433 match = np.ones(self.sites, bool)
434 for f in fields:
435 match &= np.isclose(c[f] if f in (other.coor.dtype.names or ()) else 0., self.coor[f])
436 boo += match
437 return boo
439 def plot(
440 self,
441 ms: float = 20,
442 fs: float = 20,
443 plt_index: bool = False,
444 axis: bool = False,
445 figsize: tuple[float, float] | None = None,
446 ) -> Figure:
447 '''
448 Plot the lattice in real space (its projection on the (x, y) plane,
449 in 3D space).
451 :param ms: Positive number. Default value 20. Markersize.
452 :param fs: Positive number. Default value 20. Fontsize.
453 :param plt_index: Boolean. Default value False. Plot site labels.
454 :param axis: Boolean. Default value False. Plot axis.
455 :param figsize: Tuple. Default value None. Figsize.
457 :returns:
458 * **fig** -- Figure.
459 '''
460 error_handling.empty_coor(self.coor)
461 error_handling.positive_real(ms, 'ms')
462 error_handling.positive_real(fs, 'fs')
463 error_handling.boolean(plt_index, 'plt_index')
464 error_handling.boolean(axis, 'axis')
465 if figsize is None:
466 figsize = (5, 5)
467 error_handling.list_tuple_2elem(figsize, 'figsize')
468 error_handling.positive_real(figsize[0], 'figsize[0]')
469 error_handling.positive_real(figsize[1], 'figsize[1]')
470 fig, ax = plt.subplots(figsize=figsize)
471 # plot sites
472 colors = ['b', 'r', 'g', 'y', 'm', 'k']
473 for color, tag in zip(colors, self.tags):
474 plt.plot(self.coor['x'][self.coor['tag'] == tag],
475 self.coor['y'][self.coor['tag'] == tag],
476 'o', color=color, ms=ms, markeredgecolor='none')
477 ax.set_aspect('equal')
478 ax.set_xlim(np.min(self.coor['x'])-1., np.max(self.coor['x'])+1.)
479 ax.set_ylim(np.min(self.coor['y'])-1., np.max(self.coor['y'])+1.)
480 if not axis:
481 ax.axis('off')
482 # plot indices
483 if plt_index:
484 indices = ['{}'.format(i) for i in range(self.sites)]
485 for l, x, y in zip(indices, self.coor['x'], self.coor['y']):
486 plt.annotate(l, xy=(x, y), xytext=(0, 0),
487 textcoords='offset points',
488 ha='right', va='bottom', size=fs)
489 plt.draw()
490 return fig
492 def show(self) -> None:
493 """
494 Emulate Matplotlib method plt.show().
495 """
496 plt.show()
499# Backward-compatible lowercase alias (pre-0.2 API).
500lattice = Lattice