Coverage for hiphive/force_constant_model.py: 98%
226 statements
« prev ^ index » next coverage.py v7.10.1, created at 2025-08-01 17:04 +0000
1"""
2Contains the force constant modell (FCM) which handles cluster and force
3constant information of a super cell object.
4"""
5import math
6import pickle
8from typing import IO, Tuple, Union
10import numpy as np
12from ase import Atoms
13from scipy.sparse import coo_matrix, vstack
15from .cluster_space import ClusterSpace
16from .core.atoms import Atom, atom_to_spos, spos_to_atom
17from .core.structure_alignment import align_supercell
18from .core.orbits import Orbit, OrientationFamily
19from .core.tensors import (rotation_to_cart_coord,
20 rotate_tensor_precalc, rotation_tensor_as_matrix)
21from .core.utilities import BiMap
22from .force_constants import ForceConstants, SortedForceConstants
23from .input_output.logging_tools import logger
24from .calculators.numba_calc import (clusters_force_contribution,
25 cluster_force_contribution)
27logger = logger.getChild('fcm')
30class ForceConstantModel:
31 """Transfers a cluster space onto a super structure.
33 Contains the full description of all clusters and the force constants
34 within a super cell with periodic boundary conditions.
36 Parameters
37 ----------
38 atoms
39 Configuration to which the cluster space is to be applied.
40 cs
41 A cluster space compatible with the structure of the atoms.
42 """
43 def __init__(self, atoms: Atoms, cs: ClusterSpace):
45 self.atoms = atoms.copy()
46 self.orbits = []
47 self.cluster_list = BiMap()
48 self.cs = cs
49 self._populate_orbits(atoms)
51 # TODO: refactor
52 def _populate_orbits(self, atoms: Atoms):
53 """Map the orbits from the underlying force constant potential onto the
54 supercell structure associated with this force constant model.
56 """
57 # TODO: Comment function
58 atom_lookup = {}
59 cs = self.cs
61 aligned_super_cell, scR, _ = align_supercell(
62 atoms, cs.primitive_structure, cs.symprec)
63 sc = aligned_super_cell.copy()
64 sc.cell = cs.primitive_structure.cell
65 sc.pbc = False
66 atom_list = BiMap()
67 # TODO: fix this. see also 334d8572
68 for spos in sc.get_scaled_positions():
69 atom_list.append(spos_to_atom(spos, cs.primitive_structure.basis,
70 cs.symprec))
71 sorted_cluster_list = BiMap()
72 if hasattr(cs, 'rotation_matrices'):
73 rotations = []
74 for R in cs.rotation_matrices:
75 rotations.append(rotation_to_cart_coord(
76 R, cs.primitive_structure.cell))
78 def get_atom_index(atom):
79 tupd_atom = (atom.site, *atom.offset)
80 if tupd_atom in atom_lookup:
81 return atom_lookup[tupd_atom]
82 spos = atom_to_spos(atom, cs.primitive_structure.basis)
83 pos = np.dot(spos, cs.primitive_structure.cell)
84 spos = np.dot(pos, np.linalg.inv(aligned_super_cell.cell))
85 atom = spos_to_atom(spos, aligned_super_cell.basis,
86 self.cs.symprec)
87 atom_lookup[tupd_atom] = atom.site
88 return atom.site
90 def get_mapped_cluster(cluster, offset):
91 new_cluster = []
92 for atom_index in cluster:
93 atom = cs.atom_list[atom_index]
94 translated_atom = Atom(atom.site, np.add(atom.offset, offset))
95 index = get_atom_index(translated_atom)
96 new_cluster.append(index)
97 return new_cluster
98 logger.debug('Populating orbits')
99 scR_tensormatrix_lookup = dict()
100 R_inv_tensormatrix_lookup = dict()
101 for orbit_index, orbit in enumerate(cs.orbits):
103 new_orbit = Orbit()
104 new_orbit.order = orbit.order
106 scR_tensormatrix = scR_tensormatrix_lookup.get(orbit.order, None)
107 if scR_tensormatrix is None:
108 scR_tensormatrix = rotation_tensor_as_matrix(scR, orbit.order)
109 scR_tensormatrix_lookup[orbit.order] = scR_tensormatrix
111 if len(orbit.eigentensors) > 0:
112 ets = []
113 for et in orbit.eigentensors:
114 ets.append(rotate_tensor_precalc(et, scR_tensormatrix))
115 new_orbit.eigentensors = ets
116 new_orbit.force_constant = np.zeros(ets[0].shape)
117 else:
118 new_orbit.force_constant = rotate_tensor_precalc(
119 orbit.force_constant, scR_tensormatrix)
120 cluster = cs.cluster_list[orbit.prototype_index]
121 _, pos, counts = np.unique(np.array(cluster),
122 return_index=True, return_counts=True)
123 new_orbit.positions = pos
124 prefactor = -1 / np.prod(list(map(math.factorial, counts)))
125 new_orbit.prefactors = np.array([prefactor * c for c in counts])
127 for of in orbit.orientation_families:
129 new_of = OrientationFamily()
130 if len(orbit.eigentensors) > 0:
131 R_inv_tensormatrix_index = (of.symmetry_index, orbit.order)
132 R_inv_tensormatrix = \
133 R_inv_tensormatrix_lookup.get(R_inv_tensormatrix_index,
134 None)
135 if R_inv_tensormatrix is None:
136 R_inv = rotations[of.symmetry_index].T
137 R_inv_tensormatrix = rotation_tensor_as_matrix(
138 R_inv, orbit.order)
139 R_inv_tensormatrix_lookup[R_inv_tensormatrix_index] = \
140 R_inv_tensormatrix
141 ets = []
142 for et in orbit.eigentensors:
143 et_of = rotate_tensor_precalc(et, R_inv_tensormatrix)
144 ets.append(rotate_tensor_precalc(
145 et_of, scR_tensormatrix))
146 new_of.eigentensors = ets
147 new_of.force_constant = np.zeros(ets[0].shape)
148 else:
149 new_of.force_constant = rotate_tensor_precalc(
150 of.force_constant, scR_tensormatrix)
152 cluster = cs.cluster_list[of.cluster_indices[0]]
153 if isinstance(cs, ClusterSpace):
154 perm = cs._permutations[of.permutation_indices[0]]
155 cluster = [cluster[i] for i in np.argsort(perm)]
156 for atom in atom_list:
157 if not atom.site == cs.atom_list[cluster[0]].site:
158 continue
159 offset = atom.offset
160 new_cluster = tuple(get_mapped_cluster(cluster, offset))
161 sorted_new_cluster = tuple(sorted(new_cluster))
162 if sorted_new_cluster in sorted_cluster_list: 162 ↛ 163line 162 didn't jump to line 163 because the condition on line 162 was never true
163 msg = f'Found cluster {sorted_new_cluster} in two orbits, decrease cutoff'\
164 ' (less than L/2) or increase supercell size.'
165 raise Exception(msg)
166 sorted_cluster_list.append(sorted_new_cluster)
167 self.cluster_list.append(new_cluster)
168 new_cluster_index = len(self.cluster_list) - 1
169 new_of.cluster_indices.append(new_cluster_index)
170 new_orbit.orientation_families.append(new_of)
171 self.orbits.append(new_orbit)
173 if isinstance(cs, ClusterSpace):
174 self._parameters = np.zeros(self.cs.n_dofs)
176 def get_force_constants(self) -> SortedForceConstants:
177 """Returns the force constants of the super cell.
179 Returns
180 -------
181 Complete set of force constants.
182 """
184 fc_dict = dict()
185 for orbit in self.orbits:
186 for of in orbit.orientation_families:
187 fc = of.force_constant.copy()
188 for cluster_index in of.cluster_indices:
189 cluster = self.cluster_list[cluster_index]
190 perm = np.argsort(cluster)
191 sorted_cluster = tuple(sorted(cluster))
192 fc_dict[sorted_cluster] = fc.transpose(perm)
193 return SortedForceConstants(fc_dict, self.atoms)
195 @property
196 def parameters(self) -> np.ndarray:
197 """ Parameters. """
198 return self._parameters
200 @parameters.setter
201 def parameters(self, parameters: np.ndarray) -> None:
202 self._parameters = parameters
203 mapped_parameters = self.cs._map_parameters(parameters)
204 p = 0
205 for orb in self.orbits:
206 fc_is_zero = np.allclose(orb.force_constant, 0)
207 params_is_zero = np.allclose(
208 mapped_parameters[p: p+len(orb.eigentensors)], 0)
209 if fc_is_zero and params_is_zero:
210 p += len(orb.eigentensors)
211 continue
212 orb.force_constant *= 0
213 if not params_is_zero: 213 ↛ 216line 213 didn't jump to line 216 because the condition on line 213 was always true
214 for et, a in zip(orb.eigentensors, mapped_parameters[p:]):
215 orb.force_constant += et * a
216 for of in orb.orientation_families:
217 of.force_constant *= 0
218 if not params_is_zero: 218 ↛ 216line 218 didn't jump to line 216 because the condition on line 218 was always true
219 for et, a in zip(of.eigentensors, mapped_parameters[p:]):
220 of.force_constant += et * a
221 p += len(orb.eigentensors)
223 def get_forces(self, displacements: np.ndarray) -> np.ndarray:
224 """ Returns the forces in the system given displacements.
225 The parameters of the model must be set to get any result.
227 Parameters
228 ----------
229 displacements
230 Displacements of each atom in the supercell (`N, 3` array).
231 """
232 F = np.zeros(displacements.shape)
233 f = np.zeros(3)
234 for orbit in self.orbits:
235 if np.allclose(orbit.force_constant, 0): 235 ↛ 236line 235 didn't jump to line 236 because the condition on line 235 was never true
236 continue
237 order = orbit.order
238 positions = orbit.positions
239 prefactors = orbit.prefactors
240 for of in orbit.orientation_families:
241 fc = of.force_constant.flatten()
242 fc_tmp = fc.copy()
244 for cluster_index in of.cluster_indices:
245 cluster = self.cluster_list[cluster_index]
246 cluster_force_contribution(
247 positions, prefactors, len(prefactors),
248 fc_tmp, fc, order,
249 displacements,
250 cluster, f, F)
251 return F
253 def get_fit_matrix(self, displacements: np.ndarray) -> np.ndarray:
254 """ Returns the matrix used to fit the parameters.
255 Represents the linear relation between the parameters and the forces.
257 Parameters
258 ----------
259 displacements
260 Displacements of each atom in the supercell (`(N, 3)` array).
261 """
263 F = np.zeros(displacements.shape)
264 f = np.zeros(3)
265 M = np.zeros((F.size, self.cs._cvs.shape[0]))
266 et_index = 0
267 for orbit in self.orbits:
268 fc_tmp = np.zeros(3**orbit.order)
269 for of in orbit.orientation_families:
270 clusters = np.zeros((len(of.cluster_indices), orbit.order),
271 dtype=np.int64)
272 for i, cluster_index in enumerate(of.cluster_indices):
273 clusters[i] = self.cluster_list[cluster_index]
274 for i, et in enumerate(of.eigentensors):
275 F[:] = 0
276 clusters_force_contribution(orbit.positions,
277 orbit.prefactors,
278 len(orbit.prefactors),
279 fc_tmp,
280 et.ravel(),
281 orbit.order,
282 displacements,
283 clusters, f, F)
284 M[:, et_index + i] += F.flat
285 et_index += len(orbit.eigentensors)
286 M = M.dot(self.cs._cvs.toarray())
287 return M
289 def get_energy_fit_vector(
290 self,
291 fit_matrix: np.ndarray,
292 displacements: np.ndarray,
293 ) -> np.ndarray:
294 """ Returns the energy fit vector.
295 The vector represents the linear relation between the parameters and the energy of
296 the structure.
298 Parameters
299 ----------
300 fit_matrix
301 Fit matrix for the given supercell, can be obtained by :func:`get_fit_matrix()`.
302 displacements
303 Displacements of each atom in the supercell (`(N, 3)` array).
304 """
305 energy_row = np.zeros(fit_matrix.shape[1])
306 for order in self.cs.cutoffs.orders:
307 params = self.cs.get_parameter_indices(order)
308 M = fit_matrix[:, params]
309 M = - (M.T * displacements.ravel()).T / order
310 energy_row[params] = M.sum(axis=0)
311 return energy_row
313 def get_fcs_sensing(self,
314 fcs: ForceConstants,
315 sparse: bool = False) \
316 -> Union[Tuple[np.ndarray, np.ndarray], Tuple[coo_matrix, np.ndarray]]:
317 """ Creates a fit matrix from force constants directly.
319 If the underlying cluster space can completely span the force constants
320 the correct parameters should be recovered. The method assumes that the
321 force constants obey crystal, lattice and label symmetries and will
322 naively incorporate only one force constant per orbit into the sensing
323 matrix.
325 The parameters can be extracted using, e.g., least squares from numpy::
327 parameters = np.linalg.lstsq(*fcm.get_fcs_sensing(fcs))[0]
329 Parameters
330 ----------
331 fcs
332 Force constants that are compatible with the ideal structure of the model.
334 Returns
335 -------
336 A tuple comprising the sensing matrix and the flattened, irreducible force constants.
337 """
338 M, F = [], []
339 et_index = 0
340 n_parameters = self.cs._cvs.shape[0]
341 for oi, orbit in enumerate(self.orbits):
342 cluster = self.cluster_list[orbit.prototype_index]
343 fc = fcs[cluster].ravel()
344 row, col, data = [], [], []
345 for i, et in enumerate(orbit.eigentensors):
346 for r, d in enumerate(et.flat):
347 row.append(r)
348 col.append(et_index + i)
349 data.append(d)
350 F.append(fc)
351 m = coo_matrix((data, (row, col)), shape=(3**orbit.order, n_parameters))
352 M.append(m)
353 et_index += len(orbit.eigentensors)
354 M = vstack(M)
355 M = M.dot(self.cs._cvs)
356 if not sparse:
357 M = M.toarray()
358 F = np.concatenate(F)
359 return M, F
361 @staticmethod
362 def read(f: Union[str, IO]):
363 """Reads a force constant model from file.
365 Parameters
366 ----------
367 f
368 name of input file (`str`) or stream to load from (`IO`)
370 Returns
371 -------
372 ForceConstantModel
373 Force constant model object as stored in the file.
374 """
375 if isinstance(f, str):
376 with open(f, 'rb') as fobj:
377 return pickle.load(fobj)
378 else:
379 return pickle.load(f)
381 def write(self, f: Union[str, IO]) -> None:
382 """Writes a force constant model to file.
384 Parameters
385 ----------
386 f
387 name of input file (`str`) or stream to write to (`IO`)
388 """
389 if isinstance(f, str):
390 with open(f, 'wb') as fobj:
391 pickle.dump(self, fobj)
392 else:
393 pickle.dump(self, f)