Coverage for hiphive/core/structures.py: 94%
109 statements
« prev ^ index » next coverage.py v7.10.1, created at 2025-08-01 17:04 +0000
1import numpy as np
2from ase import Atoms
5def site_offset_to_spos(
6 site: int,
7 offset: np.ndarray,
8 basis_spos: list[np.ndarray],
9) -> np.ndarray:
10 """ Returns the scaled position of an atom at the specified site and offset
11 relative to the basis in scaled coordinates.
13 Parameters
14 ----------
15 site
16 Site index.
17 offset
18 Offset in scaled coordinates (should be integers).
19 basis_spos
20 Positions of atoms in basis in scaled coordinates.
21 """
22 return offset + basis_spos[site]
25def spos_to_pos(spos: np.ndarray, cell: np.ndarray) -> np.ndarray:
26 """ Returns the Cartesian coordinate given the scaled coordinate and cell
27 metric (cell vectors as rows).
29 Parameters
30 ----------
31 spos
32 Position in scaled coordinates.
33 cell
34 Cell metric.
35 """
36 return np.dot(spos, cell)
39def pos_to_spos(pos: np.ndarray, cell: np.ndarray) -> np.ndarray:
40 """Returns the scaled coordinate given the Cartesian coordinate and the cell metric.
41 Inverse of :func:`spos_to_pos`.
43 Parameters
44 ----------
45 pos
46 Position in Cartesian coordinates.
47 cell
48 Cell metric.
49 """
50 return np.linalg.solve(cell.T, pos)
53def spos_to_site_offset(
54 spos: np.ndarray,
55 basis_spos: list[np.ndarray],
56 symprec: float,
57) -> np.ndarray:
58 """ Returns the site and offset of the atom at the specified scaled
59 coordinate given the scaled positions of the basis atoms.
61 Parameters
62 ----------
63 spos
64 Position in scaled coordinates.
65 basis_spos
66 Positions of atoms in basis in scaled coordinates.
67 symprec
68 Tolerance imposed when rounding the offset.
69 """
70 for site, sp in enumerate(basis_spos): 70 ↛ 76line 70 didn't jump to line 76 because the loop on line 70 didn't complete
71 offset = spos - sp
72 rounded_offset = offset.round(0).astype(np.int64)
73 # TODO: fix tolerance (symprec should be w.r.t. cart. coord.)
74 if np.allclose(rounded_offset, offset, rtol=0, atol=symprec):
75 return site, rounded_offset
76 raise Exception('spos {} not compatible with basis {} using symprec {}'
77 .format(spos, basis_spos, symprec))
80def pos_to_site_offset(
81 pos: np.ndarray,
82 cell: np.ndarray,
83 basis_spos: list[np.ndarray],
84 symprec: float,
85) -> np.ndarray:
86 """Returns offset given the position in Cartesian coordinates and the cell metric.
88 Parameters
89 ----------
90 pos
91 Position in Cartesian coordinates.
92 cell
93 Cell metric.
94 basis_spos
95 Positions of atoms in basis in scaled coordinates.
96 symprec
97 Tolerance imposed when rounding the offset.
98 """
99 spos = pos_to_spos(pos, cell)
100 return spos_to_site_offset(spos, basis_spos, symprec)
103def site_offset_to_pos(
104 site: int,
105 offset: np.ndarray,
106 cell: np.ndarray,
107 basis_spos: list[np.ndarray],
108) -> np.ndarray:
109 """Returns the position in Cartesian coordinates given a site index,
110 an offset, the cell metric, and the position in the basis.
112 Parameters
113 ----------
114 site
115 Site index.
116 offset
117 Offset in scaled coordinates (should be integers).
118 cell
119 Cell metric.
120 basis_spos
121 Positions of atoms in basis in scaled coordinates.
122 """
123 spos = site_offset_to_spos(site, offset, basis_spos)
124 return spos_to_pos(spos, cell)
127class BaseAtom:
128 """ This class represents an atom placed in an infinite crystal.
130 Attributes
131 ----------
132 site : int
133 Site index.
134 offset : list[int]
135 Offset in scaled coordinates.
136 """
137 def __init__(self, site, offset):
138 assert type(site) is int, type(site)
139 assert len(offset) == 3, len(offset)
140 assert (all(type(i) is int for i in offset) or
141 all(type(i) is np.int64 for i in offset)), type(offset[0])
142 self._site = site
143 self._offset = np.array(offset)
145 @property
146 def site(self) -> int:
147 """ Site index. """
148 return self._site
150 @property
151 def offset(self) -> np.ndarray:
152 """ Offset in scaled coordinates. """
153 return self._offset
155 def astype(self, dtype):
156 """ Useful arguments: list, tuple, np.int64"""
157 return dtype((self._site, *self._offset))
160class Atom(BaseAtom):
161 """ This class represents a crystal atom in a given structure.
162 """
163 def __init__(self, *args, **kwargs):
164 self._structure = kwargs.pop('structure', None)
165 super().__init__(*args, **kwargs)
167 @property
168 def pos(self) -> np.ndarray:
169 return site_offset_to_pos(self._site, self._offset,
170 self._structure.cell,
171 self._structure.spos)
173 @property
174 def number(self) -> int:
175 return self._structure.numbers[self._site]
178class SupercellAtom(Atom):
179 """ Represents an atom in a supercell but site and offset given by an
180 underlying primitve cell."""
181 def __init__(self, *args, **kwargs):
182 self._index = kwargs.pop('index')
183 assert type(self._index) is int
184 super().__init__(*args, **kwargs)
186 @property
187 def index(self):
188 return self._index
191class Structure:
192 """ This class essentially wraps the ASE :class:`Atoms` class but is a bit more
193 careful with respect to periodic boundary conditions and scaled coordinates.
194 Note that it returns :class:`hiphive.Atom` objects when queried for individual atoms.
196 Parameters
197 ----------
198 atoms
199 Atomic structure.
200 symprec
201 Tolerance imposed when rounding scaled coordinates.
202 """
203 def __init__(self, atoms: Atoms, symprec: float = 1e-6):
204 spos = atoms.get_scaled_positions(wrap=False)
205 for sp in spos.flat:
206 if not (-symprec < sp < (1 - symprec)): 206 ↛ 207line 206 didn't jump to line 207 because the condition on line 206 was never true
207 raise ValueError('bad spos {}'.format(sp))
208 self._spos = spos
209 self._cell = atoms.cell
210 self._numbers = atoms.numbers
212 def __len__(self):
213 return len(self._spos)
215 @property
216 def spos(self) -> np.ndarray:
217 """ Scaled coordinates. """
218 return self._spos
220 @property
221 def cell(self):
222 """ Cell metric. """
223 return self._cell
225 def __getitem__(self, index):
226 if index >= len(self):
227 raise IndexError('Structure contains {} atoms'.format(len(self)))
228 return Atom(index, (0, 0, 0), structure=self)
230 def atom_from_pos(self, pos: np.ndarray, symprec: float = None) -> Atom:
231 """ Returns atom given a position in Cartesian coordinates.
233 Parameters
234 ----------
235 pos
236 Position in Cartesian coordinates.
237 symprec
238 Tolerance imposed when rounding scaled coordinates.
239 """
240 if symprec is None: 240 ↛ 241line 240 didn't jump to line 241 because the condition on line 240 was never true
241 symprec = self._symprec
242 site, offset = pos_to_site_offset(pos, self._cell, self._spos, symprec)
243 return Atom(site, offset, structure=self)
246class Supercell:
247 """ This class tries to represent atoms in a supercell as positioned on the
248 primitive lattice.
250 Parameters
251 ----------
252 supercell
253 Supercell structure.
254 prim
255 Primitive structure.
256 symprec
257 Tolerance imposed when rounding scaled coordinates.
258 """
259 def __init__(self, supercell: Atoms, prim: Atoms, symprec: float):
260 self._supercell = Structure(supercell)
261 self._prim = Structure(prim)
262 self._symprec = symprec
263 self._map = list()
264 self._inverse_map_lookup = dict()
265 self._create_map()
267 def _create_map(self):
268 for atom in self._supercell:
269 atom = self._prim.atom_from_pos(atom.pos, self._symprec)
270 self._map.append(atom.astype(tuple))
272 def wrap_atom(self, atom):
273 atom = Atom(atom.site, atom.offset, structure=self._prim)
274 tup = atom.astype(tuple)
275 index = self._inverse_map_lookup.get(tup, None)
276 if index is None:
277 atom = self._supercell.atom_from_pos(atom.pos, self._symprec)
278 index = atom.site
279 self._inverse_map_lookup[tup] = index
280 return self[index]
282 def index(self, site: int, offset: np.ndarray) -> int:
283 """" Returns index of atom given a site index and an offset.
285 Parameters
286 ----------
287 site
288 Site index.
289 offset
290 Offset in scaled coordinates (should be integers).
291 """
292 atom = self.wrap_atom(BaseAtom(site, offset))
293 return atom.index
295 def __getitem__(self, index):
296 tup = self._map[index]
297 return SupercellAtom(tup[0], tup[1:], structure=self._prim,
298 index=index)
300 def __len__(self):
301 return len(self._supercell)