Coverage for tbkit/lattice.py: 100%
227 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-30 11:04 +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 Parameters
38 ----------
39 unit_cell : List of dictionaries
40 One dictionary per site within the unit cell. Each dictionary has two keys:
42 * 'tag', one-character string. Label of the associated sublattice.
43 * 'r0', Tuple. Position: :math:`(x, y)`, or :math:`(x, y, z)` for a
44 lattice in 3D space.
45 prim_vec : List of tuples
46 Define the primitive vectors. List of one/two/three tuples for 1D/2D/3D
47 lattices, each the cartesian coordinates of a primitive vector: 2-tuples
48 in the plane, or 3-tuples in 3D space (then *r0* must be 3-tuples too,
49 and *coor* gains a 'z' field, see :data:`COOR_DTYPE_3D`).
51 Notes
52 -----
53 In 3D space, the geometry methods acting in the plane (*rotation*,
54 *ellipse_in*, *boundary_line*, ...) act on :math:`(x, y)` only (rotation
55 about the z axis, cylinders, vertical planes), and *plot* draws the
56 projection on the :math:`(x, y)` plane.
58 Examples
59 --------
60 Example usage::
62 # Line-Centered Square lattice
63 unit_cell = [{'tag': 'a', 'r0': (0., 0.)}, {'tag': 'a', 'r0': (0., 1.)}]
64 prim_vec = [(0, 2), (2, 0)]
65 lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec)
66 '''
68 def __init__(self, unit_cell: list[dict], prim_vec: Sequence[tuple[float, ...]]) -> None:
69 error_handling.unit_cell(unit_cell)
70 error_handling.prim_vec(prim_vec)
71 error_handling.space_dim(unit_cell, prim_vec)
72 self.unit_cell = unit_cell
73 self.prim_vec = prim_vec
74 self.tags = np.unique(np.array([dic['tag'] for dic in self.unit_cell]))
75 self.space_dim = len(prim_vec[0]) # 2, or 3 for a lattice in 3D space
76 self.dtype = COOR_DTYPE if self.space_dim == 2 else COOR_DTYPE_3D
77 self.n1, self.n2, self.n3 = 0, 0, 0
78 self.coor = np.array([], dtype=self.dtype)
79 self.sites = 0
81 def get_lattice(self, n1: int, n2: int = 1, n3: int = 1) -> None:
82 r'''
83 Get the lattice positions, sorted by (y, x) -- by (z, y, x) in 3D space.
85 Parameters
86 ----------
87 n1 : Positive Integer
88 Number of unit cells along :math:`\mathbf{a}_1`.
89 n2 : Positive Integer, default 1
90 Number of unit cells along :math:`\mathbf{a}_2`.
91 n3 : Positive Integer, default 1
92 Number of unit cells along :math:`\mathbf{a}_3` (3D lattices only).
94 Examples
95 --------
96 Example usage::
98 # Line-Centered Square lattice
99 unit_cell = [{'tag': 'a', 'r0': (0., 0.)}, {'tag': 'a', 'r0': (0., 1.)}]
100 prim_vec = [(0, 2), (2, 0)]
101 lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec)
102 lat.get_lattice(n1=4, n2=5)
103 '''
104 error_handling.get_lattice(self.prim_vec, n1, n2, n3)
105 sites_uc = len(self.unit_cell)
106 sites_tag = n1*n2*n3
107 self.sites = sites_uc * sites_tag
108 self.coor = np.empty(self.sites, dtype=self.dtype)
109 self.n1, self.n2, self.n3 = n1, n2, n3
110 # translations R = i1 a1 + i2 a2 (+ i3 a3), i1 running fastest
111 axes = [f for f in ('x', 'y', 'z')[:self.space_dim]]
112 trans = np.zeros((sites_tag, self.space_dim))
113 counts = (n1, n2, n3)
114 index = np.indices(counts[:max(len(self.prim_vec), 1)][::-1]).reshape(
115 max(len(self.prim_vec), 1), -1)[::-1]
116 for n, a in zip(index, self.prim_vec):
117 trans += n[:, None] * np.array(a, dtype='f8')[None, :]
118 for i, dic in enumerate(self.unit_cell):
119 for c, f in enumerate(axes):
120 self.coor[f][i*sites_tag: (i+1)*sites_tag] = trans[:, c] + dic['r0'][c]
121 self.coor['tag'][i*sites_tag: (i+1)*sites_tag] = dic['tag']
122 self.coor = np.sort(self.coor, order=self.sort_order())
124 def sort_order(self) -> tuple[str, ...]:
125 '''
126 Private method. Fields the sites are sorted by: ('y', 'x'), or
127 ('z', 'y', 'x') in 3D space.
128 '''
129 return ('y', 'x') if self.space_dim == 2 else ('z', 'y', 'x')
131 def distances(self) -> tuple[NDArray, ...]:
132 '''
133 Private method. Pairwise coordinate differences, *d[i, j]* from site
134 *i* to site *j*: (dx, dy), or (dx, dy, dz) in 3D space.
135 '''
136 return tuple(self.coor[f] - self.coor[f].reshape(self.sites, 1)
137 for f in ('x', 'y', 'z')[:self.space_dim])
139 def add_sites(self, coor: NDArray) -> None:
140 '''
141 Add sites.
143 Parameters
144 ----------
145 coor
146 Structured array with keys: {'x', 'y', 'tag'}
147 ({'x', 'y', 'z', 'tag'} for a lattice in 3D space).
149 Examples
150 --------
151 Example usage::
153 # Square lattice
154 unit_cell = [{'tag': 'a', 'r0': (0., 0.)}]
155 prim_vec = [(0, 1), (1, 0)]
156 lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec)
157 lat.get_lattice(n1=2, n2=2)
158 coor = np.array([(-1., -1, 'b'), (-2., -2, 'c')],
159 dtype=[('x', 'f8'), ('y', 'f8'), ('tag', 'U1')])
160 lat.add_sites(coor)
161 '''
162 error_handling.coor(coor, self.dtype)
163 self.coor = np.concatenate([self.coor, coor])
164 self.sites += len(coor)
165 self.tags = np.unique(np.concatenate([self.tags, coor['tag']]))
166 self.coor = np.sort(self.coor, order=self.sort_order())
168 def remove_sites(self, index: list[int]) -> None:
169 '''
170 Remove sites defined by their indices
171 (use method lattice.plot(plt_index=True)
172 to get access to the site indices).
174 Parameters
175 ----------
176 index : List
177 Site indices to be removed.
179 Examples
180 --------
181 Example usage::
183 # Square lattice
184 unit_cell = [{'tag': 'a', 'r0': (0., 0.)}]
185 prim_vec = [(0, 1), (1, 0)]
186 lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec)
187 lat.get_lattice(n1=2, n2=2)
188 lat.remove_sites([0, 2])
189 '''
190 error_handling.empty_coor(self.coor)
191 error_handling.remove_sites(index, self.sites)
192 mask = np.ones(self.sites, bool)
193 mask[index] = False
194 self.coor = self.coor[mask]
195 self.sites = self.coor.size
197 def remove_dangling(self) -> None:
198 '''
199 Remove dangling sites
200 (sites connected with just another site).
201 '''
202 error_handling.empty_coor(self.coor)
203 while True:
204 coords = np.stack([self.coor[f] for f in ('x', 'y', 'z')[:self.space_dim]], axis=1)
205 tree = cKDTree(coords)
206 near = tree.query(coords, k=min(self.sites, 8))[0][:, 1:]
207 len_hop = np.min(near[near > 0]) # the shortest distance between sites
208 pairs = tree.query_pairs(len_hop * (1. + 1e-5) + 1e-8, output_type='ndarray')
209 dis = np.linalg.norm(coords[pairs[:, 1]] - coords[pairs[:, 0]], axis=1)
210 pairs = pairs[np.isclose(dis, len_hop)]
211 degree = np.bincount(pairs.ravel(), minlength=self.sites)
212 dang = list(np.flatnonzero(degree == 1))
213 self.coor = np.delete(self.coor, dang, axis=0)
214 self.sites -= len(dang)
215 if dang == []:
216 break
218 def shift_x(self, shift: float) -> None:
219 '''
220 Shift the x coordinates.
222 Parameters
223 ----------
224 shift : Real number
225 Shift value.
226 '''
227 error_handling.empty_coor(self.coor)
228 error_handling.real_number(shift, 'shift')
229 self.coor['x'] += shift
231 def shift_y(self, shift: float) -> None:
232 '''
233 Shift the y coordinates.
235 Parameters
236 ----------
237 shift : Real number
238 Shift value.
239 '''
240 error_handling.empty_coor(self.coor)
241 error_handling.real_number(shift, 'shift')
242 self.coor['y'] += shift
244 def change_sign_x(self) -> None:
245 '''
246 Change x coordinates sign.
247 '''
248 error_handling.empty_coor(self.coor)
249 self.coor['x'] *= -1
251 def change_sign_y(self) -> None:
252 '''
253 Change y coordinates sign.
254 '''
255 error_handling.empty_coor(self.coor)
256 self.coor['y'] *= -1
258 def shift_z(self, shift: float) -> None:
259 '''
260 Shift the z coordinates (3D lattices only).
262 Parameters
263 ----------
264 shift : Real number
265 Shift value.
266 '''
267 error_handling.empty_coor(self.coor)
268 error_handling.space_3d(self.space_dim)
269 error_handling.real_number(shift, 'shift')
270 self.coor['z'] += shift
272 def change_sign_z(self) -> None:
273 '''
274 Change z coordinates sign (3D lattices only).
275 '''
276 error_handling.empty_coor(self.coor)
277 error_handling.space_3d(self.space_dim)
278 self.coor['z'] *= -1
280 def slab(self, z_min: float, z_max: float) -> None:
281 r'''
282 Keep only the sites with :math:`z_{min} < z < z_{max}` (3D lattices only).
284 Parameters
285 ----------
286 z_min : Real number
287 z_max : Real number, larger than *z_min*
288 '''
289 error_handling.empty_coor(self.coor)
290 error_handling.space_3d(self.space_dim)
291 error_handling.real_number(z_min, 'z_min')
292 error_handling.real_number(z_max, 'z_max')
293 error_handling.smaller(z_min, 'z_min', z_max, 'z_max')
294 self.coor = self.coor[(self.coor['z'] > z_min) & (self.coor['z'] < z_max)]
295 self.sites = len(self.coor)
297 def boundary_line(self, cx: float, cy: float, co: float) -> None:
298 r'''
299 Keep only the sites with :math:`c_yy+c_xx > c_0`.
301 Parameters
302 ----------
303 cx : Real number
304 :math:`c_x` value.
305 cy : Real number
306 :math:`c_y` value.
307 co : Real number
308 :math:`c_0` value.
309 '''
310 error_handling.empty_coor(self.coor)
311 error_handling.real_number(cx, 'cx')
312 error_handling.real_number(cy, 'cy')
313 error_handling.real_number(co, 'co')
314 self.coor = self.coor[cy * self.coor['y'] + cx * self.coor['x'] > co]
315 self.sites = len(self.coor)
317 def ellipse_in(self, rx: float, ry: float, x0: float, y0: float) -> None:
318 r'''
319 Select sites according to
321 .. math::
323 (x-x_0)^2/r_x^2+(y-y_0)^2/r_y^2 < 1\, .
325 Parameters
326 ----------
327 rx : Positive Real number
328 Radius along :math:`x`.
329 ry : Positive Real number
330 Radius along :math:`y`.
331 x0 : Real number
332 :math:`x` center.
333 y0 : Real number
334 :math:`y` center.
335 '''
336 error_handling.empty_coor(self.coor)
337 error_handling.positive_real(rx, 'rx')
338 error_handling.positive_real(ry, 'ry')
339 error_handling.real_number(x0, 'x0')
340 error_handling.real_number(y0, 'y0')
341 self.coor = self.coor[(self.coor['x'] -x0) ** 2 / rx ** 2 + \
342 (self.coor['y'] -y0) ** 2 / ry ** 2 < 1.]
343 self.sites = len(self.coor)
345 def ellipse_out(self, rx: float, ry: float, x0: float, y0: float) -> None:
346 r'''
347 Select sites according to
349 .. math::
351 (x-x_0)^2/r_x^2+(y-y_0)^2/r_y^2 > 1\, .
353 Parameters
354 ----------
355 rx : Positive Real number
356 Radius along :math:`x`.
357 ry : Positive Real number
358 Radius along :math:`y`.
359 x0 : Real number
360 :math:`x` center.
361 y0 : Real number
362 :math:`y` center.
363 '''
364 error_handling.empty_coor(self.coor)
365 error_handling.positive_real(rx, 'rx')
366 error_handling.positive_real(ry, 'ry')
367 error_handling.real_number(x0, 'x0')
368 error_handling.real_number(y0, 'y0')
369 self.coor = self.coor[(self.coor['x'] -x0) ** 2 / rx ** 2 + \
370 (self.coor['y'] -y0) ** 2 / ry ** 2 > 1.]
371 self.sites = len(self.coor)
373 def center(self) -> None:
374 '''
375 Fix the center of mass of the lattice at (0, 0) -- (0, 0, 0) in 3D space.
376 '''
377 error_handling.empty_coor(self.coor)
378 for f in ('x', 'y', 'z')[:self.space_dim]:
379 self.coor[f] -= np.mean(self.coor[f])
381 def rotation(self, theta: float) -> None:
382 r'''
383 Rotate the lattice structure about the origin (about the z axis, in
384 3D space) by the angle :math:`\theta`.
386 Parameters
387 ----------
388 theta
389 Rotation angle in degrees.
390 '''
391 error_handling.empty_coor(self.coor)
392 error_handling.real_number(theta, 'theta')
393 theta *= PI / 180
394 x = self.coor['x'].copy()
395 y = self.coor['y'].copy()
396 self.coor['x'] = x * np.cos(theta) - y * np.sin(theta)
397 self.coor['y'] = y * np.cos(theta) + x * np.sin(theta)
399 def clean_coor(self) -> None:
400 '''
401 Keep only the sites with different coordinates.
402 '''
403 error_handling.empty_coor(self.coor)
404 fields = list(('x', 'y', 'z')[:self.space_dim])
405 coor = self.coor[fields].copy()
406 for f in fields:
407 coor[f] = self.coor[f].round(4)
408 _, idx = np.unique(coor, return_index=True)
409 self.coor = self.coor[idx]
410 self.sites = len(self.coor)
412 def __add__(self, other: 'Lattice') -> 'Lattice':
413 '''
414 Overloading operator +.
415 '''
416 error_handling.lat(other)
417 error_handling.empty_coor(self.coor)
418 error_handling.empty_coor(other.coor)
419 coor = np.concatenate([self.coor, other.coor])
420 tags = np.concatenate([self.tags, other.tags])
421 lat = lattice(unit_cell=self.unit_cell, prim_vec=self.prim_vec)
422 lat.add_sites(coor)
423 lat.sites = self.sites + other.sites
424 lat.tags = np.unique(tags)
425 return lat
427 def __iadd__(self, other: 'Lattice') -> 'Lattice':
428 '''
429 Overloading operator +=.
430 '''
431 error_handling.lat(other)
432 error_handling.empty_coor(self.coor)
433 error_handling.empty_coor(other.coor)
434 self.coor = np.concatenate([self.coor, other.coor])
435 self.sites += other.sites
436 self.tags = np.unique(np.concatenate([self.tags, other.tags]))
437 return self
439 def __sub__(self, other: 'Lattice') -> 'Lattice':
440 '''
441 Overloading operator -.
443 .. note::
445 The tags are not considered in the lattice subtraction.
446 '''
447 error_handling.lat(other)
448 error_handling.empty_coor(self.coor)
449 error_handling.empty_coor(other.coor)
450 boo = self.overlaps(other)
451 coor = self.coor[np.logical_not(boo)]
452 lat = lattice(unit_cell=self.unit_cell, prim_vec=self.prim_vec)
453 lat.add_sites(coor)
454 lat.sites = len(lat.coor)
455 lat.tags = self.tags
456 return lat
458 def __isub__(self, other: 'Lattice') -> 'Lattice':
459 '''
460 Overloading operator -=.
462 .. note::
464 The tags are not considered in the lattice subtraction.
465 '''
466 error_handling.lat(other)
467 error_handling.empty_coor(self.coor)
468 error_handling.empty_coor(other.coor)
469 boo = self.overlaps(other)
470 self.coor = self.coor[np.logical_not(boo)]
471 self.sites = int(np.sum(np.logical_not(boo)))
472 return self
474 def overlaps(self, other: 'Lattice') -> NDArray[np.bool_]:
475 '''
476 Private method. Mask of the sites of *self* that coincide with a site
477 of *other*.
478 '''
479 names = self.coor.dtype.names or ()
480 fields = [f for f in ('x', 'y', 'z') if f in names]
481 boo = np.zeros(self.sites, bool)
482 for c in other.coor:
483 match = np.ones(self.sites, bool)
484 for f in fields:
485 match &= np.isclose(c[f] if f in (other.coor.dtype.names or ()) else 0., self.coor[f])
486 boo += match
487 return boo
489 def plot(
490 self,
491 ms: float = 20,
492 fs: float = 20,
493 plt_index: bool = False,
494 axis: bool = False,
495 figsize: tuple[float, float] | None = None,
496 ) -> Figure:
497 '''
498 Plot the lattice in real space (its projection on the (x, y) plane,
499 in 3D space).
501 Parameters
502 ----------
503 ms : Positive number, default 20
504 Markersize.
505 fs : Positive number, default 20
506 Fontsize.
507 plt_index : Boolean, default False
508 Plot site labels.
509 axis : Boolean, default False
510 Plot axis.
511 figsize : Tuple, default None
512 Figsize.
514 Returns
515 -------
516 fig : Figure
517 '''
518 error_handling.empty_coor(self.coor)
519 error_handling.positive_real(ms, 'ms')
520 error_handling.positive_real(fs, 'fs')
521 error_handling.boolean(plt_index, 'plt_index')
522 error_handling.boolean(axis, 'axis')
523 if figsize is None:
524 figsize = (5, 5)
525 error_handling.list_tuple_2elem(figsize, 'figsize')
526 error_handling.positive_real(figsize[0], 'figsize[0]')
527 error_handling.positive_real(figsize[1], 'figsize[1]')
528 fig, ax = plt.subplots(figsize=figsize)
529 # plot sites
530 colors = ['b', 'r', 'g', 'y', 'm', 'k']
531 for color, tag in zip(colors, self.tags):
532 plt.plot(self.coor['x'][self.coor['tag'] == tag],
533 self.coor['y'][self.coor['tag'] == tag],
534 'o', color=color, ms=ms, markeredgecolor='none')
535 ax.set_aspect('equal')
536 ax.set_xlim(np.min(self.coor['x'])-1., np.max(self.coor['x'])+1.)
537 ax.set_ylim(np.min(self.coor['y'])-1., np.max(self.coor['y'])+1.)
538 if not axis:
539 ax.axis('off')
540 # plot indices
541 if plt_index:
542 indices = ['{}'.format(i) for i in range(self.sites)]
543 for l, x, y in zip(indices, self.coor['x'], self.coor['y']):
544 plt.annotate(l, xy=(x, y), xytext=(0, 0),
545 textcoords='offset points',
546 ha='right', va='bottom', size=fs)
547 plt.draw()
548 return fig
550 def show(self) -> None:
551 """
552 Emulate Matplotlib method plt.show().
553 """
554 plt.show()
557# Backward-compatible lowercase alias (pre-0.2 API).
558lattice = Lattice