Coverage for hiphive/force_constants.py: 95%
290 statements
« prev ^ index » next coverage.py v7.10.1, created at 2025-08-01 17:04 +0000
1"""
2This module provides functionality for storing and handling of force constants.
3"""
5import numpy as np
6import tarfile
8from abc import ABC, abstractmethod
9from itertools import product
10from string import ascii_lowercase, ascii_uppercase
11from typing import List, Tuple
13from scipy.linalg import eig
14from ase import Atoms, units
16from .input_output import shengBTE as shengBTEIO
17from .input_output import phonopy as phonopyIO
18from .input_output.read_write_files import add_items_to_tarfile_pickle, read_items_pickle, \
19 add_ase_atoms_to_tarfile, read_ase_atoms_from_tarfile
20from .input_output import gpumd as gpumdIO
23class ForceConstants(ABC):
24 """ Base class for force constants. """
26 def __init__(self, supercell: Atoms):
27 if not all(supercell.pbc):
28 raise ValueError('Configuration must have periodic boundary conditions.')
29 self._supercell = supercell.copy()
31 @abstractmethod
32 def __getitem__(self):
33 pass
35 @property
36 def n_atoms(self) -> int:
37 """ Number of atoms. """
38 return len(self.supercell)
40 @property
41 def supercell(self) -> Atoms:
42 """ Supercell associated with force constants. """
43 return self._supercell.copy()
45 @property
46 def clusters(self) -> list:
47 """ Sorted list of clusters """
48 return sorted(self._fc_dict.keys(), key=lambda c: (len(c), c))
50 def __len__(self) -> int:
51 """ Number of clusters (or force constants). """
52 return len(self._fc_dict)
54 def __add__(self, fcs_other) -> None:
55 """ Addition of two force constants """
56 if not isinstance(fcs_other, type(self)):
57 raise ValueError('ForceConstants objects are of different types')
59 # if they have overlapping orders, raise
60 if any(order in self.orders for order in fcs_other.orders):
61 raise ValueError('ForceConstants objects share at least one order')
63 # check that they share the same supercell
64 if self.supercell != fcs_other.supercell:
65 raise ValueError('supercells do not match')
67 # add force constants
68 fc_merged = {**self._fc_dict, **fcs_other._fc_dict}
69 return self.__class__(fc_merged, self.supercell)
71 def get_fc_dict(self, order: int = None) -> dict:
72 """ Returns force constant dictionary for one specific order.
74 The returned dict may be sparse or may be dense depending on the
75 underlying force constants.
77 Parameters
78 ----------
79 order
80 Force constants returned for this order.
82 Returns
83 -------
84 Dictionary with keys corresponding to clusters and values to
85 respective force constant.
86 """
87 if order is None:
88 return self._fc_dict
90 if order not in self.orders:
91 raise ValueError('Order {} not in ForceConstants'.format(order))
93 fc_order = {}
94 for c, fc in self._fc_dict.items():
95 if len(c) == order:
96 fc_order[c] = fc
97 return fc_order
99 def get_fc_array(self, order: int, format: str = 'phonopy') -> np.ndarray:
100 """ Returns force constants in array format for specified order.
102 Parameters
103 ----------
104 order
105 Force constants for this order will be returned.
106 format
107 Specify which format (shape) the NumPy array should have,
108 possible values are `'phonopy'` and `'ase'`.
110 Returns
111 -------
112 Array with shape `(N,)*order + (3,)*order` where `N` is
113 the number of atoms.
114 """
115 if order not in self.orders:
116 raise ValueError('Order not in orders.')
117 if format not in ['ase', 'phonopy']:
118 raise ValueError('Format must be either ase or phonopy.')
120 # generate fc array for phonopy format
121 fc_array = np.zeros((self.n_atoms, ) * order + (3, ) * order)
122 for cluster in product(range(self.n_atoms), repeat=order):
123 fc_array[cluster] = self[cluster]
125 if format == 'ase':
126 if order != 2: 126 ↛ 127line 126 didn't jump to line 127 because the condition on line 126 was never true
127 raise ValueError('ASE format works only for order 2')
128 return fc_array.transpose([0, 2, 1, 3]).reshape(
129 self.n_atoms * 3, self.n_atoms * 3)
130 else:
131 return fc_array
133 def compute_gamma_frequencies(self) -> np.ndarray:
134 """ Returns the Gamma frequencies in THz using the second-order force
135 constants. """
136 fc2_array = self.get_fc_array(order=2)
137 masses = self.supercell.get_masses()
138 return _compute_gamma_frequencies(fc2_array, masses)
140 def assert_acoustic_sum_rules(self, order: int = None, tol: float = 1e-6):
141 """ Asserts that force constants obey acoustic sum rules.
143 Parameters
144 ----------
145 order
146 Specifies which order to check, if None all are checked.
147 tol
148 Numeric tolerance for checking sum rules.
150 Raises
151 ------
152 AssertionError
153 If the acoustic sum rules are violated.
154 """
156 # set up orders
157 if order is None:
158 orders = self.orders
159 else:
160 if order not in self.orders:
161 raise ValueError('Order not available in FCS')
162 orders = [order]
164 atomic_indices = range(self.n_atoms)
165 for order in orders:
166 assert_msg = 'Acoustic sum rule for order {} violated for atom'
167 assert_msg += ' {}' * (order - 1) + ' x'
168 for ijk in product(atomic_indices, repeat=order-1):
169 fc_sum_ijk = np.zeros((3, )*order)
170 for m in atomic_indices:
171 cluster = ijk + (m, )
172 fc_sum_ijk += self[cluster]
173 assert np.linalg.norm(fc_sum_ijk) < tol, assert_msg.format(order, *ijk)
175 def print_force_constant(self, cluster: Tuple[int]) -> None:
176 """
177 Prints force constants for a cluster in a nice format.
179 Parameters
180 ----------
181 cluster
182 Sites belonging to the cluster.
183 """
184 print(self._repr_fc(cluster))
186 def __eq__(self, other):
188 # check supercells are the same
189 if not len(self.supercell) == len(other.supercell):
190 return False
191 if not np.allclose(self.supercell.positions, other.supercell.positions):
192 return False
193 if not np.allclose(self.supercell.cell, other.supercell.cell):
194 return False
195 if not all(self.supercell.numbers == other.supercell.numbers):
196 return False
198 # check orders and clusters are the same
199 if self.orders != other.orders:
200 return False
201 if self.clusters != other.clusters: 201 ↛ 202line 201 didn't jump to line 202 because the condition on line 201 was never true
202 return False
204 # check force constants
205 for c in self.clusters:
206 if not np.allclose(self[c], other[c]):
207 return False
208 return True
210 def __str__(self) -> str:
211 s = []
212 s.append(' ForceConstants '.center(54, '='))
213 s.append('Orders: {}'.format(self.orders))
214 s.append('Atoms: {}'.format(self.supercell))
215 s.append('')
216 if len(self) > 10: 216 ↛ 223line 216 didn't jump to line 223 because the condition on line 216 was always true
217 for c in self.clusters[:3]:
218 s.append(self._repr_fc(c)+'\n')
219 s.append('...\n')
220 for c in self.clusters[-3:]:
221 s.append(self._repr_fc(c)+'\n')
222 else:
223 for c in self.clusters:
224 s.append(self._repr_fc(c)+'\n')
225 return '\n'.join(s)
227 def __repr__(self) -> str:
228 fc_dict_str = '{{{}: {}, ..., {}: {}}}'.format(
229 self.clusters[0], self[self.clusters[0]].round(5),
230 self.clusters[-1], self[self.clusters[-1]].round(5))
231 return ('ForceConstants(fc_dict={}, atoms={!r})'
232 .format(fc_dict_str, self.supercell))
234 def _repr_fc(self, cluster: Tuple[int],
235 precision: float = 5, suppress: bool = True) -> str:
236 """
237 Representation for single cluster and its force constant.
239 Parameters
240 ----------
241 cluster
242 Tuple of ints indicating the sites belonging to the cluster.
243 """
244 s = []
245 s.append('Cluster: {}'.format(cluster))
246 for atom_index in cluster:
247 s.append(str(self.supercell[atom_index]))
248 s.append('Force constant:')
249 s.append(np.array_str(self[cluster], precision=precision,
250 suppress_small=suppress))
251 return '\n'.join(s)
253 def _sanity_check_dict(self, fc_dict: dict) -> None:
254 """ Checks that all indices in clusters are between 0 and number of
255 atoms. """
256 # Use flat numpy array for fast check
257 cluster_indices = np.concatenate([cluster for cluster in fc_dict.keys()])
258 if not np.all((0 <= cluster_indices) & (cluster_indices < self.n_atoms)):
259 # If the check failed, do slower list processing to find the erring cluster
260 for cluster in fc_dict.keys(): 260 ↛ exitline 260 didn't return from function '_sanity_check_dict' because the loop on line 260 didn't complete
261 if not all(0 <= i < self.n_atoms for i in cluster):
262 raise ValueError('Cluster {} not in supercell'.format(cluster))
264 @classmethod
265 def from_arrays(cls, supercell: Atoms,
266 fc2_array: np.ndarray = None, fc3_array: np.ndarray = None):
267 """ Constructs FCs from numpy arrays.
268 One or both of :attr:`fc2_array` and :attr:`fc3_array` must not be `None`.
270 Parameters
271 ----------
272 supercell
273 Aupercell structure.
274 fc2_array
275 Second-order force constant in phonopy format, i.e., in shape `(N, N, 3, 3)`.
276 fc3_array
277 Third-order force constant in phonopy format, i.e., in shape `(N, N, N, 3, 3, 3)`.
278 """
279 if fc2_array is None and fc3_array is None:
280 raise ValueError('Please provide force constant arrays.')
282 n_atoms = len(supercell)
283 if fc2_array is None:
284 fc2_dict = dict()
285 else:
286 if fc2_array.shape != (n_atoms, n_atoms, 3, 3):
287 raise ValueError('fc2_array has wrong shape.')
288 fc2_dict = array_to_dense_dict(fc2_array)
290 if fc3_array is None:
291 fc3_dict = dict()
292 else:
293 if fc3_array.shape != (n_atoms, n_atoms, n_atoms, 3, 3, 3):
294 raise ValueError('fc3_array has wrong shape.')
295 fc3_dict = array_to_dense_dict(fc3_array)
297 fc_dict = {**fc2_dict, **fc3_dict}
298 return cls.from_dense_dict(fc_dict, supercell)
300 @classmethod
301 def from_sparse_dict(cls, fc_dict: dict, supercell: Atoms):
302 """ Assumes label symmetries, meaning only one cluster for each
303 permuation should be included.
305 Parameters
306 ----------
307 fc_dict
308 Force constants with keys corresponding to clusters
309 and values to the force constants.
310 supercell
311 Atomic configuration.
312 """
313 return SortedForceConstants(fc_dict, supercell=supercell)
315 @classmethod
316 def from_dense_dict(cls, fc_dict: dict, supercell: Atoms):
317 """ All permutations of clusters that are not zero must be listed,
318 if label symmetries are fullfilled will return a :class:`SortedForceConstants`.
320 Parameters
321 ----------
322 fc_dict
323 Force constants with keys corresponding to clusters
324 and values to the force constants.
325 supercell
326 Atomic configuration.
327 """
328 if check_label_symmetries(fc_dict):
329 fc_sparse_dict = dense_dict_to_sparse_dict(fc_dict)
330 return SortedForceConstants(fc_sparse_dict, supercell=supercell)
331 else:
332 return RawForceConstants(fc_dict, supercell=supercell)
334 @classmethod
335 def read_phonopy(cls, supercell: Atoms, fname: str, format: str = None):
336 """ Reads force constants from a phonopy calculation.
338 Parameters
339 ----------
340 supercell
341 Supercell structure (corresponding to `SPOSCAR`).
342 fname
343 Name of second-order force constant file.
344 format
345 Format for second-order force constants;
346 possible values: `'text'`, `'hdf5'`.
347 """
348 fc2_array = phonopyIO.read_phonopy_fc2(fname, format=format)
349 return cls.from_arrays(supercell, fc2_array)
351 @classmethod
352 def read_phono3py(cls, supercell: Atoms, fname: str):
353 """ Reads force constants from a phono3py calculation.
355 Parameters
356 ----------
357 supercell
358 Supercell structure (corresponding to `SPOSCAR`).
359 fname
360 Name of third-order force constant file.
361 """
362 fc3_array = phonopyIO.read_phonopy_fc3(fname)
363 return cls.from_arrays(supercell, fc3_array=fc3_array)
365 @classmethod
366 def read_shengBTE(cls, supercell: Atoms, fname: str, prim: Atoms, symprec=1e-5):
367 """ Reads third order force constants from a shengBTE calculation.
368 The shengBTE force constants will be mapped onto a supercell.
370 Parameters
371 ----------
372 supercell
373 Supercell structure.
374 fname
375 Name of third-order force constant file.
376 prim
377 Primitive configuration (must be equivalent to structure used in
378 the shengBTE calculation).
379 symprec
380 Structural symmetry tolerance.
381 """
382 fcs = shengBTEIO.read_shengBTE_fc3(fname, prim, supercell, symprec)
383 return fcs
385 @classmethod
386 def read(cls, fname: str):
387 """ Reads a :class:`ForceConstants` object from file.
389 Parameters
390 ----------
391 fname
392 Name of file from which to read.
393 """
394 tar_file = tarfile.open(mode='r', name=fname)
395 items = read_items_pickle(tar_file, 'fc_dict')
396 fc_dict = items['fc_dict']
397 fcs_type = items['fcs_type']
398 supercell = read_ase_atoms_from_tarfile(tar_file, 'supercell')
399 tar_file.close()
400 if fcs_type == 'SortedForceConstants':
401 return SortedForceConstants(fc_dict, supercell)
402 elif fcs_type == 'RawForceConstants': 402 ↛ 405line 402 didn't jump to line 405 because the condition on line 402 was always true
403 return RawForceConstants(fc_dict, supercell)
404 else:
405 raise ValueError('FCS type not recongnized')
407 def write(self, fname: str) -> None:
408 """ Writes :class:`ForceConstants` object to file.
410 Parameters
411 ----------
412 fname
413 Name of file to which to write.
414 """
415 tar_file = tarfile.open(name=fname, mode='w')
416 items_pickle = dict(fc_dict=self._fc_dict, fcs_type=self.__class__.__name__)
417 add_items_to_tarfile_pickle(tar_file, items_pickle, 'fc_dict')
418 add_ase_atoms_to_tarfile(tar_file, self.supercell, 'supercell')
419 tar_file.close()
421 def write_to_phonopy(self, fname: str, format: str = None) -> None:
422 """
423 Writes force constants in phonopy format.
425 Parameters
426 ----------
427 fname
428 Name of file to which to write second-order force constant.
429 format
430 Format for second-order force constants;
431 possible values: `'text'`, `'hdf5'`.
432 """
433 phonopyIO.write_phonopy_fc2(fname, self, format=format)
435 def write_to_phono3py(self, fname: str) -> None:
436 """
437 Writes force constants in phono3py format.
439 Parameters
440 ----------
441 fname
442 Name of file to which to write third-order force constant.
443 """
444 phonopyIO.write_phonopy_fc3(fname, self)
446 def write_to_shengBTE(self, fname: str, prim: Atoms, **kwargs) -> None:
447 """
448 Writes third-order force constants in shengBTE format.
450 Parameters
451 ----------
452 fname
453 Name of file to which to write third-order force constant.
454 prim
455 Primitive configuration (must be equivalent to structure used in
456 the shengBTE calculation).
457 """
458 shengBTEIO.write_shengBTE_fc3(fname, self, prim, **kwargs)
461class SortedForceConstants(ForceConstants):
462 """ Force constants with label symmetries.
464 Parameters
465 ----------
466 fc_dict : dict
467 Force constants with keys corresponding to clusters and values to
468 the force constants; should only contain sorted clusters.
469 supercell
470 Supercell structure.
471 """
473 def __init__(self, fc_dict: dict, supercell: Atoms) -> None:
474 super().__init__(supercell)
475 self._sanity_check_dict(fc_dict)
476 self._fc_dict = fc_dict
477 self._orders = sorted(set(len(c) for c in self._fc_dict.keys()))
479 def __getitem__(self, cluster: Tuple[int]):
480 sorted_cluster = tuple(sorted(cluster))
482 # cluster not in fcs
483 if sorted_cluster not in self._fc_dict.keys():
484 return np.zeros((3,)*len(cluster))
486 # return fc for the unsorted cluster
487 inv_perm = np.argsort(np.argsort(cluster))
488 return self._fc_dict[sorted_cluster].transpose(inv_perm)
490 @property
491 def orders(self) -> List[int]:
492 """ Orders for which force constants exist. """
493 return self._orders.copy()
495 def _sanity_check_dict(self, fc_dict: dict) -> None:
496 super()._sanity_check_dict(fc_dict)
498 # also check clusters are sorted
499 for cluster in fc_dict.keys():
500 if cluster != tuple(sorted(cluster)):
501 raise ValueError('Found unsorted cluster {}'.format(cluster))
503 def write_to_GPUMD(
504 self,
505 fname_fc: str,
506 fname_clusters: str,
507 order: int,
508 tol: float = 1e-10,
509 ):
510 """
511 Writes force constants of the specified order in GPUMD format.
513 Parameters
514 ----------
515 fname_fc
516 Name of file which contains the look-up force constants.
517 fname_clusters
518 Name of file which contains the clusters and the fc lookup index.
519 order
520 Force constants for this order will be written to file.
521 tol
522 If the norm of a force constant is less than tol then it is not written.
523 If two force-constants are within tol; they are considered equal.
524 """
525 gpumdIO.write_fcs_gpumd(fname_fc, fname_clusters, self, order, tol)
528class RawForceConstants(ForceConstants):
529 """ Force constants without label symmetries.
531 Parameters
532 ----------
533 fc_dict : dict
534 Force constants with keys corresponding to clusters and values to the force constants,
535 should contain all clusters with nonzero force constants.
536 supercell
537 Supercell structure.
538 """
540 def __init__(self, fc_dict: dict, supercell: Atoms) -> None:
541 super().__init__(supercell)
542 self._sanity_check_dict(fc_dict)
543 self._fc_dict = fc_dict
544 self._orders = sorted(set(len(c) for c in self._fc_dict.keys()))
546 def __getitem__(self, cluster: Tuple[int]):
547 if cluster not in self._fc_dict.keys():
548 return np.zeros((3,)*len(cluster))
549 return self._fc_dict[cluster]
551 @property
552 def orders(self) -> List[int]:
553 """ Orders for which force constants exist. """
554 return self._orders.copy()
557# ====================== #
558# HELPER FUNCTIONS #
559# ====================== #
562def array_to_dense_dict(fc_array: np.ndarray, fc_tol: float = 1e-10) -> dict:
563 """ Constructs a dense dict from a FC array in phonopy format.
565 Force constants with norm smaller than :attr:`fc_tol` will be considered zero and
566 not included in the output dictionary.
568 Parameters
569 ----------
570 fc_array
571 Force constant array in phonopy format.
572 fc_tol
573 Force constants below this value are considered zero and omitted from the output dict.
574 """
576 # sanity check
577 n_atoms = fc_array.shape[0]
578 order = int(len(fc_array.shape) / 2)
579 if fc_array.shape != (n_atoms, ) * order + (3, ) * order: 579 ↛ 580line 579 didn't jump to line 580 because the condition on line 579 was never true
580 raise ValueError('fc_array has the wrong shape.')
582 # construct dense dict
583 fc_dict = dict()
584 for cluster in product(range(n_atoms), repeat=order):
585 fc = fc_array[cluster]
586 if np.linalg.norm(fc) > fc_tol:
587 fc_dict[cluster] = fc
588 return fc_dict
591def check_label_symmetries(fc_dict: dict) -> bool:
592 """ Checks label symmetries for dense FC dict.
594 Parameters
595 ----------
596 fc_dict
597 Force constants with keys corresponding to clusters and values to the force constants.
598 """
599 # TODO: tol, which one to use etc
600 for cluster, fc in fc_dict.items():
601 inv_perm = np.argsort(np.argsort(cluster))
602 sorted_cluster = tuple(sorted(cluster))
603 if sorted_cluster not in fc_dict: 603 ↛ 604line 603 didn't jump to line 604 because the condition on line 603 was never true
604 return False
605 if not np.allclose(fc, fc_dict[sorted_cluster].transpose(inv_perm)):
606 return False
607 return True
610def dense_dict_to_sparse_dict(fc_dict: dict) -> dict:
611 """ Converts dense dict to sparse dict.
613 This does not check if label symmetry is True, but rather will just keep
614 the sorted clusters and their force constants.
616 Parameters
617 ----------
618 fc_dict
619 Force constants with keys corresponding to clusters and values to the force constants.
620 """
621 fc_dict_sparse = dict()
622 for cluster, fc in fc_dict.items():
623 if cluster == tuple(sorted(cluster)):
624 fc_dict_sparse[cluster] = fc
625 return fc_dict_sparse
628def symbolize_force_constant(fc: np.ndarray, tol: float = 1e-10) -> np.ndarray:
629 """Carries out a symbolic symmetrization of a force constant tensor.
631 Parameters
632 ----------
633 fc
634 Force constant tensor.
635 tol
636 Tolerance used to decide whether two elements are identical.
638 Returns
639 -------
640 Symbolic representation of force constant matrix.
641 """
642 fc_int = np.round(fc / tol).astype(int)
643 fc_chars = np.empty(fc_int.shape, dtype=object).flatten()
644 all_chars = ascii_lowercase + ascii_uppercase
645 lookup_chars = {}
646 for i, val in enumerate(fc_int.flatten()):
647 if val == 0:
648 fc_chars[i] = 0
649 elif val in lookup_chars.keys():
650 fc_chars[i] = lookup_chars[val]
651 elif -val in lookup_chars.keys():
652 fc_chars[i] = '-{:}'.format(lookup_chars[-val])
653 else:
654 lookup_chars[val] = all_chars[len(lookup_chars.keys())]
655 fc_chars[i] = lookup_chars[val]
656 return fc_chars.reshape(fc_int.shape)
659def _compute_gamma_frequencies(fc2: np.ndarray, masses: List[float]) -> np.ndarray:
660 """ Computes Gamma frequencies using second-order force constants.
661 Assumes :attr:`fc2` is in units of eV/A2.
663 Parameters
664 ----------
665 fc2
666 Decond-order force constants in phonopy format.
667 masses
668 Mass of each atom.
670 Returns
671 -------
672 Gamma frequencies in THz.
673 """
675 n_atoms = fc2.shape[0]
676 if len(masses) != n_atoms: 676 ↛ 677line 676 didn't jump to line 677 because the condition on line 676 was never true
677 raise ValueError('Length of masses not compatible with fc2')
678 mass_matrix = np.sqrt(np.outer(masses, masses))
680 # divide with mass matrix
681 fc2_tmp = np.zeros((n_atoms, n_atoms, 3, 3))
682 for pair in product(range(n_atoms), repeat=2):
683 fc2_tmp[pair] = fc2[pair] / mass_matrix[pair]
685 # reshape into matrix and solve eigenvalues
686 fc2_tmp = fc2_tmp.transpose([0, 2, 1, 3]).reshape(n_atoms * 3, n_atoms * 3)
687 eigen_vals, _ = eig(fc2_tmp)
688 eigen_vals *= 1e20 / units.J * units.kg # [1/s**2]
689 eigen_vals.sort()
691 # if negative eigenval, set frequency to negative
692 gamma_frequencies = []
693 for val in eigen_vals.real:
694 if val >= 0:
695 gamma_frequencies.append(np.sqrt(val))
696 else:
697 gamma_frequencies.append(-np.sqrt(np.abs(val)))
699 # Sort and convert to THz
700 gamma_frequencies = np.array(gamma_frequencies) / 1e12 / (2 * np.pi)
701 gamma_frequencies.sort()
702 return gamma_frequencies