Coverage for hiphive/force_constants.py: 95%
290 statements
« prev ^ index » next coverage.py v7.6.8, created at 2024-11-28 11:20 +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 """ int : number of clusters (or force constants) """
52 return len(self._fc_dict)
54 def __add__(self, fcs_other) -> None:
55 """ ForceConstants : 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 NumPy 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 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.
269 One or both of fc2_array and fc3_array must not be None
271 Parameters
272 ----------
273 supercell
274 supercell structure
275 fc2_array
276 second-order force constant in phonopy format, i.e. must have shape (N, N, 3, 3)
277 fc3_array
278 third-order force constant in phonopy format, i.e. must have shape (N, N, N, 3, 3, 3)
279 """
280 if fc2_array is None and fc3_array is None:
281 raise ValueError('Please provide force constant arrays.')
283 n_atoms = len(supercell)
284 if fc2_array is None:
285 fc2_dict = dict()
286 else:
287 if fc2_array.shape != (n_atoms, n_atoms, 3, 3):
288 raise ValueError('fc2 array has wrong shape')
289 fc2_dict = array_to_dense_dict(fc2_array)
291 if fc3_array is None:
292 fc3_dict = dict()
293 else:
294 if fc3_array.shape != (n_atoms, n_atoms, n_atoms, 3, 3, 3):
295 raise ValueError('fc2 array has wrong shape')
296 fc3_dict = array_to_dense_dict(fc3_array)
298 fc_dict = {**fc2_dict, **fc3_dict}
299 return cls.from_dense_dict(fc_dict, supercell)
301 @classmethod
302 def from_sparse_dict(cls, fc_dict: dict, supercell: Atoms):
303 """ Assumes label symmetries, meaning only one cluster for each
304 permuation should be included
306 Parameters
307 ----------
308 fc_dict
309 keys corresponding to clusters 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 SortedForceConstants
320 Parameters
321 ----------
322 fc_dict
323 keys corresponding to clusters and values to the force constants
324 supercell
325 atomic configuration
326 """
327 if check_label_symmetries(fc_dict):
328 fc_sparse_dict = dense_dict_to_sparse_dict(fc_dict)
329 return SortedForceConstants(fc_sparse_dict, supercell=supercell)
330 else:
331 return RawForceConstants(fc_dict, supercell=supercell)
333 @classmethod
334 def read_phonopy(cls, supercell: Atoms, fname: str, format: str = None):
335 """ Reads force constants from a phonopy calculation.
337 Parameters
338 ----------
339 supercell
340 supercell structure (`SPOSCAR`)
341 fname
342 name of second-order force constant file
343 format
344 format for second-order force constants;
345 possible values: "text", "hdf5"
346 """
347 fc2_array = phonopyIO.read_phonopy_fc2(fname, format=format)
348 return cls.from_arrays(supercell, fc2_array)
350 @classmethod
351 def read_phono3py(cls, supercell: Atoms, fname: str):
352 """ Reads force constants from a phono3py calculation.
354 Parameters
355 ----------
356 supercell
357 supercell structure (`SPOSCAR`)
358 fname
359 name of third-order force constant file
360 """
361 fc3_array = phonopyIO.read_phonopy_fc3(fname)
362 return cls.from_arrays(supercell, fc3_array=fc3_array)
364 @classmethod
365 def read_shengBTE(cls, supercell: Atoms, fname: str, prim: Atoms, symprec=1e-5):
366 """ Reads third order force constants from a shengBTE calculation.
368 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 ForceConstants 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 entire 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 keys corresponding to clusters and values to the force constants,
468 should only contain sorted clusters
469 supercell : ase.Atoms
470 """
472 def __init__(self, fc_dict: dict, supercell: Atoms) -> None:
473 super().__init__(supercell)
474 self._sanity_check_dict(fc_dict)
475 self._fc_dict = fc_dict
476 self._orders = sorted(set(len(c) for c in self._fc_dict.keys()))
478 def __getitem__(self, cluster: Tuple[int]):
479 sorted_cluster = tuple(sorted(cluster))
481 # cluster not in fcs
482 if sorted_cluster not in self._fc_dict.keys():
483 return np.zeros((3,)*len(cluster))
485 # return fc for the unsorted cluster
486 inv_perm = np.argsort(np.argsort(cluster))
487 return self._fc_dict[sorted_cluster].transpose(inv_perm)
489 @property
490 def orders(self) -> List[int]:
491 """ orders for which force constants exist """
492 return self._orders.copy()
494 def _sanity_check_dict(self, fc_dict: dict) -> None:
495 super()._sanity_check_dict(fc_dict)
497 # also check clusters are sorted
498 for cluster in fc_dict.keys():
499 if cluster != tuple(sorted(cluster)):
500 raise ValueError('Found unsorted cluster {}'.format(cluster))
502 def write_to_GPUMD(self, fname_fc, fname_clusters, order, tol=1e-10):
503 """
504 Writes force constants of the specified order in GPUMD format.
506 Parameters
507 ----------
508 fname_fc : str
509 name of file which contains the lookup force constants
510 fname_clusters : str
511 name of file which contains the clusters and the fc lookup index
512 order : int
513 force constants for this order will be written to file
514 tol : float
515 if the norm of a force constant is less than tol then it is not written.
516 if two force-constants are within tol; they are considered equal.
517 """
518 gpumdIO.write_fcs_gpumd(fname_fc, fname_clusters, self, order, tol)
521class RawForceConstants(ForceConstants):
522 """ Force constants without label symmetries.
524 Parameters
525 ----------
526 fc_dict : dict
527 keys corresponding to clusters and values to the force constants,
528 should contain all clusters with nonzero force constants
529 supercell : ase.Atoms
530 """
532 def __init__(self, fc_dict: dict, supercell: Atoms) -> None:
533 super().__init__(supercell)
534 self._sanity_check_dict(fc_dict)
535 self._fc_dict = fc_dict
536 self._orders = sorted(set(len(c) for c in self._fc_dict.keys()))
538 def __getitem__(self, cluster: Tuple[int]):
539 if cluster not in self._fc_dict.keys():
540 return np.zeros((3,)*len(cluster))
541 return self._fc_dict[cluster]
543 @property
544 def orders(self) -> List[int]:
545 """ orders for which force constants exist """
546 return self._orders.copy()
549# ====================== #
550# HELPER FUNCTIONS #
551# ====================== #
554def array_to_dense_dict(fc_array: np.ndarray, fc_tol: float = 1e-10) -> dict:
555 """ Constructs a dense dict from an fc array in phonopy format.
557 Force constants with norm smaller than fc_tol will be considered zero and
558 therefore not included in the fc_dict.
560 Parameters
561 ----------
562 fc_array
563 force constant array in phonopy format
564 fc_tol
565 tolerance for considering force constants zero or not
566 """
568 # sanity check
569 n_atoms = fc_array.shape[0]
570 order = int(len(fc_array.shape) / 2)
571 if fc_array.shape != (n_atoms, ) * order + (3, ) * order: 571 ↛ 572line 571 didn't jump to line 572 because the condition on line 571 was never true
572 raise ValueError('fc array has bad shape')
574 # construct dense dict
575 fc_dict = dict()
576 for cluster in product(range(n_atoms), repeat=order):
577 fc = fc_array[cluster]
578 if np.linalg.norm(fc) > fc_tol:
579 fc_dict[cluster] = fc
580 return fc_dict
583def check_label_symmetries(fc_dict: dict) -> bool:
584 """ Checks label symmetries for dense fc dict.
586 TODO
587 ----
588 tol, which one to use etc
590 Parameters
591 ----------
592 fc_dict
593 keys corresponding to clusters and values to the force constants
594 """
595 for cluster, fc in fc_dict.items():
596 inv_perm = np.argsort(np.argsort(cluster))
597 sorted_cluster = tuple(sorted(cluster))
598 if sorted_cluster not in fc_dict: 598 ↛ 599line 598 didn't jump to line 599 because the condition on line 598 was never true
599 return False
600 if not np.allclose(fc, fc_dict[sorted_cluster].transpose(inv_perm)):
601 return False
602 return True
605def dense_dict_to_sparse_dict(fc_dict: dict) -> dict:
606 """ Converts dense dict to sparse dict.
608 This does not check if label symmetry is True, but rather will just keep
609 the sorted clusters and their force constants.
611 Parameters
612 ----------
613 fc_dict
614 keys corresponding to clusters and values to the force constants
615 """
616 fc_dict_sparse = dict()
617 for cluster, fc in fc_dict.items():
618 if cluster == tuple(sorted(cluster)):
619 fc_dict_sparse[cluster] = fc
620 return fc_dict_sparse
623def symbolize_force_constant(fc: np.ndarray, tol: float = 1e-10) -> np.ndarray:
624 """Carries out a symbolic symmetrization of a force constant tensor.
626 Parameters
627 ----------
628 fc
629 force constant tensor
630 tol
631 tolerance used to decide whether two elements are identical
633 Returns
634 -------
635 symbolic representation of force constant matrix
636 """
637 fc_int = np.round(fc / tol).astype(int)
638 fc_chars = np.empty(fc_int.shape, dtype=object).flatten()
639 all_chars = ascii_lowercase + ascii_uppercase
640 lookup_chars = {}
641 for i, val in enumerate(fc_int.flatten()):
642 if val == 0:
643 fc_chars[i] = 0
644 elif val in lookup_chars.keys():
645 fc_chars[i] = lookup_chars[val]
646 elif -val in lookup_chars.keys():
647 fc_chars[i] = '-{:}'.format(lookup_chars[-val])
648 else:
649 lookup_chars[val] = all_chars[len(lookup_chars.keys())]
650 fc_chars[i] = lookup_chars[val]
651 return fc_chars.reshape(fc_int.shape)
654def _compute_gamma_frequencies(fc2: np.ndarray, masses: List[float]) -> np.ndarray:
655 """ Computes Gamma frequencies using second-order force constants.
656 Assumes fc2 is in units of eV/A2.
658 Parameters
659 ----------
660 fc2
661 second-order force constants in phonopy format
662 masses
663 mass of each atom
665 Returns
666 -------
667 Gamma frequencies in THz
668 """
670 n_atoms = fc2.shape[0]
671 if len(masses) != n_atoms: 671 ↛ 672line 671 didn't jump to line 672 because the condition on line 671 was never true
672 raise ValueError('Length of masses not compatible with fc2')
673 mass_matrix = np.sqrt(np.outer(masses, masses))
675 # divide with mass matrix
676 fc2_tmp = np.zeros((n_atoms, n_atoms, 3, 3))
677 for pair in product(range(n_atoms), repeat=2):
678 fc2_tmp[pair] = fc2[pair] / mass_matrix[pair]
680 # reshape into matrix and solve eigenvalues
681 fc2_tmp = fc2_tmp.transpose([0, 2, 1, 3]).reshape(n_atoms * 3, n_atoms * 3)
682 eigen_vals, _ = eig(fc2_tmp)
683 eigen_vals *= 1e20 / units.J * units.kg # [1/s**2]
684 eigen_vals.sort()
686 # if negative eigenval, set frequency to negative
687 gamma_frequencies = []
688 for val in eigen_vals.real:
689 if val >= 0:
690 gamma_frequencies.append(np.sqrt(val))
691 else:
692 gamma_frequencies.append(-np.sqrt(np.abs(val)))
694 # Sort and convert to THz
695 gamma_frequencies = np.array(gamma_frequencies) / 1e12 / (2 * np.pi)
696 gamma_frequencies.sort()
697 return gamma_frequencies