Coverage for dynasor / modes / atoms.py: 99%
82 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-16 12:31 +0000
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-16 12:31 +0000
1from typing import Union
2import numpy as np
3from ase.units import fs
4from ase import Atoms
5from numpy.typing import NDArray
6from .tools import inv
7from ..units import Dalton_to_dmu
10class DynasorAtoms:
11 """Dynasor's representation of a structure."""
12 def __init__(self, atoms: Atoms):
13 """Initialized using an ASE Atoms object"""
14 self._atoms = atoms
16 @property
17 def pos(self) -> NDArray[float]:
18 """Cartesian positions."""
19 return self._atoms.positions.copy()
21 @property
22 def positions(self) -> NDArray[float]:
23 """Cartesian positions."""
24 return self.pos
26 @property
27 def spos(self) -> NDArray[float]:
28 """Reduced (or scaled) positions of atoms."""
29 return self._atoms.get_scaled_positions()
31 @property
32 def scaled_positions(self) -> NDArray[float]:
33 """Reduced (or scaled) positions of atoms."""
34 return self.spos
36 @property
37 def cell(self) -> NDArray[float]:
38 """Cell of atoms with cell vectors as rows."""
39 return self._atoms.cell.array.copy()
41 @property
42 def inv_cell(self) -> NDArray[float]:
43 """The inverse cell transpose so the inverse cell vectors are rows, no 2pi."""
44 return np.linalg.inv(self._atoms.cell.array).T
46 @property
47 def numbers(self) -> NDArray[int]:
48 """Chemical number for each atom, e.g., 1 for H, 2 for He etc."""
49 return self._atoms.numbers.copy()
51 @property
52 def masses(self) -> NDArray[float]:
53 """Masses of atoms in dmu."""
54 return self._atoms.get_masses() / fs ** 2 # In eVfs²/Ų
56 @property
57 def volume(self) -> float:
58 """Volume of cell."""
59 return self._atoms.get_volume()
61 @property
62 def n_atoms(self) -> int:
63 """Number of atoms."""
64 return len(self._atoms)
66 @property
67 def symbols(self) -> list[str]:
68 """List of chemical symbol for each element."""
69 return list(self._atoms.symbols)
71 def to_ase(self) -> Atoms:
72 """Converts the internal Atoms to ASE :class:`Atoms`."""
73 return Atoms(cell=self.cell, numbers=self.numbers, positions=self.positions, pbc=True)
75 def __repr__(self) -> str:
76 return str(self)
79class Prim(DynasorAtoms):
80 def __str__(self):
81 strings = [f"""Primitive cell:
82Number of atoms: {self.n_atoms}
83Volume: {self.volume:.3f}
84Atomic species present: {set(self.symbols)}
85Atomic numbers present: {set([int(n) for n in self.numbers])}
86Cell:
87[[{self.cell[0, 0]:<20}, {self.cell[0, 1]:<20}, {self.cell[0, 2]:<20}],
88 [{self.cell[1, 0]:<20}, {self.cell[1, 1]:<20}, {self.cell[1, 2]:<20}],
89 [{self.cell[2, 0]:<20}, {self.cell[2, 1]:<20}, {self.cell[2, 2]:<20}]]
90"""]
91 strings.append(f"{'Ind':<5}{'Sym':<5}{'Num':<5}{'Mass (Da)':<10}{'x':<10}{'y':<10}{'z':<10}"
92 f"{'a':<10}{'b':<10}{'c':<10}")
93 atom_s = []
94 for i, p, sp, m, n, s in zip(
95 range(self.n_atoms), self.positions, self.spos, self.masses / Dalton_to_dmu,
96 self.numbers, [a.symbol for a in self.to_ase()]):
97 atom_s.append(f'{i:<5}{s:<5}{n:<5}{m:<10.2f}{p[0]:<10.3f}{p[1]:<10.3f}{p[2]:<10.3f}'
98 f'{sp[0]:<10.3f}{sp[1]:<10.3f}{sp[2]:<10.3f}')
100 strings = strings + atom_s
102 string = '\n'.join(strings)
104 return string
107class Supercell(DynasorAtoms):
108 """The supercell takes care of some mappings between the primitive and repeated structure.
110 In particular the P-matrix connecting the cells as well as the offset-index of each atom is
111 calculated.
113 Note that the positions cannot be revovered as `offset x cell + basis` since the atoms get
114 wrapped.
116 Parameters
117 ----------
118 supercell
119 Some ideal repetition of the primitive structure and possible wrapping.
120 prim
121 Primitive structure.
122 """
124 def __init__(self, supercell: Union[Atoms, DynasorAtoms], prim: Union[Atoms, DynasorAtoms]):
125 self.prim = Prim(prim.copy())
126 super().__init__(supercell)
128 # determine P-matrix relating supercell to primitive cell
129 from dynasor.tools.structures import get_P_matrix
130 self._P = get_P_matrix(self.prim.cell, self.cell) # P C = S
131 self._P_inv = inv(self.P)
133 # find the index and offsets for supercell using primitive as base unit
134 from dynasor.tools.structures import get_offset_index
135 self._offsets, self._indices = get_offset_index(prim, supercell, wrap=True)
137 @property
138 def P(self) -> NDArray[float]:
139 """P-matrix is defined as dot(P, prim.cell) = supercell.cell"""
140 return self._P.copy()
142 @property
143 def P_inv(self) -> NDArray[float]:
144 """Inverse of `P`."""
145 return self._P_inv.copy()
147 @property
148 def offsets(self) -> NDArray[float]:
149 """The offset of each atom."""
150 return self._offsets.copy()
152 @property
153 def indices(self) -> NDArray[int]:
154 """The basis index of each atom"""
155 return self._indices.copy()
157 @property
158 def n_cells(self) -> int:
159 """Number of unit cells"""
160 return self.n_atoms // self.prim.n_atoms
162 def __str__(self):
164 string = f"""Supercell:
165Number of atoms: {self.n_atoms}
166Volume: {self.volume:.3f}
167Number of unit cells: {self.n_cells}
168Cell:
169[[{self.cell[0, 0]:<20}, {self.cell[0, 1]:<20}, {self.cell[0, 2]:<20}],
170 [{self.cell[1, 0]:<20}, {self.cell[1, 1]:<20}, {self.cell[1, 2]:<20}],
171 [{self.cell[2, 0]:<20}, {self.cell[2, 1]:<20}, {self.cell[2, 2]:<20}]]
172P-matrix:
173[[{self.P[0, 0]:<20}, {self.P[0, 1]:<20}, {self.P[0, 2]:<20}],
174 [{self.P[1, 0]:<20}, {self.P[1, 1]:<20}, {self.P[1, 2]:<20}],
175 [{self.P[2, 0]:<20}, {self.P[2, 1]:<20}, {self.P[2, 2]:<20}]]
176{self.prim}
177"""
178 return string