Coverage for hiphive/core/atoms.py: 89%
72 statements
« prev ^ index » next coverage.py v7.6.8, created at 2024-11-28 11:20 +0000
« prev ^ index » next coverage.py v7.6.8, created at 2024-11-28 11:20 +0000
1"""
2Collection of functions and classes for handling information concerning atoms
3and structures, including the relationship between primitive cell and
4supercells that are derived thereof.
5"""
7import pickle
8from ase import Atoms as aseAtoms
9import numpy as np
10from ..input_output.logging_tools import logger
12# TODO: Rename logger
13logger = logger.getChild('atoms')
16class Atom:
17 # TODO: This class should inherit some immutable to make it clear that
18 # there is no reference to any other obj
19 """Unique representation of an atom in a lattice with a basis
21 Class for storing information about the position of an atom in a supercell
22 relative to the origin of the underlying primitive cell. This class is used
23 for handling the relationship between a primitive cell and supercells
24 derived thereof.
26 Parameters
27 ----------
28 site : int
29 site index
30 offset : list(float) or numpy.ndarray
31 must contain three elements, offset_x, offset_y, offset_z
32 """
33 def __init__(self, site, offset):
34 offset = tuple(offset)
35 self._site = site
36 self._offset = offset
38 @property
39 def site(self):
40 """int : index of corresponding site in the primitive basis"""
41 return self._site
43 @property
44 def offset(self):
45 """list(int) : translational offset of the supercell site relative
46 to the origin of the primitive cell in units of primitive lattice
47 vectors"""
48 return self._offset
50 def __repr__(self):
51 return 'Atom({}, {})'.format(self.site, self.offset)
53 def spos(atom, basis):
54 return np.add(basis[atom.site], atom.offset)
56 def pos(atom, basis, cell):
57 spos = atom.spos(basis)
58 return np.dot(spos, cell)
60 @staticmethod
61 def spos_to_atom(spos, basis, tol=None):
62 # TODO: Why is this simply duplicated spos_to_atom from helper below?
63 if not tol: 63 ↛ 65line 63 didn't jump to line 65 because the condition on line 63 was never true
64 # TODO: Link to config file
65 tol = 1e-4
66 for site, base in enumerate(basis): 66 ↛ 76line 66 didn't jump to line 76 because the loop on line 66 didn't complete
67 offset = np.subtract(spos, base)
68 diff = offset - np.round(offset, 0)
69 if np.linalg.norm(diff) < tol:
70 offset = np.round(offset, 0).astype(int)
71 atom = Atom(site, offset)
72 assert np.linalg.norm(spos - atom.spos(basis)) < tol, (
73 '{} with basis {} != {}'.format(atom, basis, spos))
74 return atom
76 s = '{} not compatible with {} and tolerance {}'
77 raise Exception(s.format(spos, basis, tol))
79 def __hash__(self):
80 return hash((self._site, *self.offset))
82 def __eq__(self, other):
83 if not isinstance(other, Atom): 83 ↛ 84line 83 didn't jump to line 84 because the condition on line 83 was never true
84 return False
85 return self.site == other.site and self.offset == other.offset
88class Atoms(aseAtoms):
89 """Minimally augmented version of the ASE Atoms class suitable for handling
90 primitive cell information.
92 Saves and loads by pickle.
93 """
94 @property
95 def basis(self):
96 """numpy.ndarray : scaled coordinates of the sites in the primitive basis
97 """
98 return self.get_scaled_positions().round(12) % 1
100 def write(self, f):
101 """ Writes the object to file.
103 Note: Only the cell, basis and numbers are stored!
105 Parameters
106 ----------
107 f : str or file object
108 name of input file (str) or stream to write to (file object)
109 """
110 data = {}
111 data['cell'] = self.cell
112 data['basis'] = self.basis
113 data['numbers'] = self.numbers
115 pickle.dump(data, f)
117 @staticmethod
118 def read(f):
119 """ Load an hiPhive Atoms object from file.
121 Parameters
122 ----------
123 f : str or file object
124 name of input file (str) or stream to load from (file object)
126 Returns
127 -------
128 hiPhive Atoms object
129 """
130 data = pickle.load(f)
131 atoms = aseAtoms(numbers=data['numbers'],
132 scaled_positions=data['basis'],
133 cell=data['cell'],
134 pbc=True)
135 return Atoms(atoms)
138def atom_to_spos(atom, basis):
139 """Helper function for obtaining the position of a supercell atom in scaled
140 coordinates.
142 Parameters
143 ----------
144 atom : hiPhive.Atom
145 supercell atom
146 basis : list(list(float)) or numpy.ndarray
147 positions of sites in the primitive basis
149 Returns
150 -------
151 numpy.ndarray
152 scaled coordinates of an atom in a supercell
153 """
154 return np.add(atom.offset, basis[atom.site])
157def spos_to_atom(spos, basis, tol=1e-4):
158 """Helper function for transforming a supercell position to the primitive
159 basis.
161 Parameters
162 ----------
163 spos : list(list(float)) or numpy.ndarray
164 scaled coordinates of an atom in a supercell
165 basis : list(list(float)) or numpy.ndarray
166 positions of sites in the primitive basis
167 tol : float
168 a general tolerance
170 Returns
171 -------
172 hiphive.Atom
173 supercell atom
174 """
175 # TODO: Fix tolerance
176 # If needed, convert inputs to arrays to make use of numpy vectorization
177 spos = np.asarray(spos)
178 basis = np.asarray(basis)
179 # If the scaled position belongs to this site, the offset is the
180 # difference in scaled coordinates and should be integer
181 offsets = spos - basis
182 # The diff is the difference between the offset vector and the nearest
183 # integer vector.
184 diffs = offsets - np.round(offsets, 0)
185 # It should be close to the null vector if this is the correct site.
186 match_indices = np.nonzero(np.linalg.norm(diffs, axis=1) < tol)[0]
187 # If no atom was found or more than one atoms were found we throw an error
188 if len(match_indices) != 1: 188 ↛ 189line 188 didn't jump to line 189 because the condition on line 188 was never true
189 raise ValueError(f'{spos} not compatible with {basis} and tolerance {tol}')
191 # This should be the correct atom
192 site = match_indices[0]
193 # If the difference is less than the tol make the offset integers
194 offset = np.rint(offsets[site])
195 atom = Atom(site, offset)
196 # Just to be sure we check that the atom actually produces the
197 # input spos given the input basis
198 s = ('Atom=[{},{}] with basis {} != {}'
199 .format(atom.site, atom.offset, basis, spos))
200 assert np.linalg.norm(spos - atom_to_spos(atom, basis)) < tol, s
201 return atom