refactor

phonopy_vibspec/phonons_analyzer.py

@@ -1,17 +1,19 @@
 import pathlib
-import yaml
 import numpy
+
 import phonopy
-from phonopy.interface import calculator
+from phonopy.interface import calculator as phonopy_calculator
 
 from numpy.typing import NDArray
 from typing import Optional, List
 
+from phonopy_vibspec.spectra import RamanSpectrum, InfraredSpectrum
+
 THZ_TO_INV_CM = 33.35641
 
 # [(value, coef), ...]
 # see https://en.wikipedia.org/wiki/Finite_difference_coefficient
-TWO_POINTS_STENCIL = [[-1, -.5], [1, .5]]  # two-points, centered
+TWO_POINTS_STENCIL = [(-1, -.5), (1, .5)]  # two-points, centered
 
 
 class PhononsAnalyzer:
@@ -37,7 +39,7 @@ def __init__(self, phonon: phonopy.Phonopy):
         # get irreps
         self.phonotopy.set_irreps([0, 0, 0])
         self.irreps = phonon.get_irreps()
-        self.irrep_labels = [None] * (self.N * 3)
+        self.irrep_labels = [''] * (self.N * 3)
 
         # TODO: that's internal API, so subject to change!
         for label, dgset in zip(self.irreps._get_ir_labels(), self.irreps._degenerate_sets):
@@ -61,6 +63,25 @@ def from_phonopy(
             born_filename=born_filename,
         ))
 
+    def infrared_spectrum(self, modes: Optional[List[int]] = None) -> InfraredSpectrum:
+        """
+        The `modes` is a 0-based list of mode to include.
+        If `modes` is None, then all non-acoustic modes are selected.
+        """
+        born_tensor = self.phonotopy.nac_params['born']
+
+        # select modes if any
+        if modes is None:
+            modes = list(range(3, 3 * self.N))
+
+        frequencies = [self.frequencies[m] for m in modes]
+        irrep_labels = [self.irrep_labels[m] for m in modes]
+        disps = self.eigendisps[numpy.ix_(modes)]
+
+        dmu_dq = numpy.einsum('ijb,jab->ia', disps, born_tensor)
+
+        return InfraredSpectrum(modes=modes, frequencies=frequencies, irrep_labels=irrep_labels, dmu_dq=dmu_dq)
+
     def infrared_intensities(self, selected_modes: Optional[List[int]] = None) -> NDArray[float]:
         """Compute the infrared intensities with Eq. 7 of 10.1039/C7CP01680H.
         Intensities are given in [e²/AMU].
@@ -76,34 +97,34 @@ def infrared_intensities(self, selected_modes: Optional[List[int]] = None) -> ND
         dipoles = numpy.einsum('ijb,jab->ia', disps, born_tensor)
         return (dipoles ** 2).sum(axis=1)
 
-    def create_displaced_geometries(
+    def prepare_raman(
         self,
         directory: pathlib.Path,
-        disp: float = 1e-2,
         modes: Optional[List[int]] = None,
+        disp: float = 1e-2,
         ref: Optional[str] = None,
-        stencil: List[List[float]] = TWO_POINTS_STENCIL
-    ):
+        stencil: Optional[list] = None
+    ) -> RamanSpectrum:
         """
-        Create a set of displaced geometries of the unitcell in `directory`.
-        The number of geometries that are generated depends on the stencil.
+        Prepare a Raman spectrum by creating a set of displaced geometries of the unitcell in `directory`.
+        The number of geometries that are generated per mode depends on the stencil.
 
-        The `modes` is a 0-based list of mode.
-        If `mode` is `None`,  all mode are selected, except the acoustic ones, and only one version of degenerated ones.
+        The `modes` is a 0-based list of mode to include.
+        If `modes` is None, then all non-acoustic modes are selected.
         """
 
-        # select modes
+        stencil = TWO_POINTS_STENCIL if stencil is None else stencil
+
+        # select modes if any
         if modes is None:
-            modes = []
-            for dgset in self.irreps._degenerate_sets:
-                if any(x < 3 for x in dgset):  # skip acoustic phonons
-                    continue
+            modes = list(range(3, 3 * self.N))
 
-                modes.append(dgset[0])
+        frequencies = [self.frequencies[m] for m in modes]
+        irrep_labels = [self.irrep_labels[m] for m in modes]
 
         # create geometries
+        steps = []
         base_geometry = self.phonotopy.unitcell
-        raman_disps_info = []
 
         for mode in modes:
             if mode < 0 or mode >= 3 * self.N:
@@ -113,19 +134,26 @@ def create_displaced_geometries(
             if ref == 'norm':
                 step = disp / float(numpy.linalg.norm(self.eigendisps[mode]))
 
-            raman_disps_info.append({'mode': mode, 'step': step})
+            steps.append(step)
 
             for i, (value, _) in enumerate(stencil):
                 displaced_geometry = base_geometry.copy()
                 displaced_geometry.set_positions(
                     base_geometry.positions + value * step * self.eigendisps[mode]
                 )
 
-                calculator.write_crystal_structure(
+                phonopy_calculator.write_crystal_structure(
                     directory / 'unitcell_{:04d}_{:02d}.vasp'.format(mode + 1, i + 1),  # 1-based output
                     displaced_geometry,
                     interface_mode='vasp'
                 )
 
-        with (directory / 'raman_disps.yml').open('w') as f:
-            yaml.dump({'stencil': stencil, 'modes': raman_disps_info}, f)
+        calculator = RamanSpectrum(
+            modes=modes,
+            frequencies=frequencies,
+            irrep_labels=irrep_labels,
+            steps=steps,
+            stencil=stencil
+        )
+
+        return calculator

phonopy_vibspec/scripts/displaced_geometries.py

@@ -42,13 +42,15 @@ def main():
         force_constants_filename=args.fc,
     )
 
-    phonons.create_displaced_geometries(
-        pathlib.Path('.'),
+    raman_spectrum = phonons.prepare_raman(
+        args.output,
         disp=args.displacement,
         ref=args.ref,
         modes=args.modes if len(args.modes) > 0 else None
     )
 
+    raman_spectrum.to_hdf5(args.output / 'raman.hdf5')
+
 
 if __name__ == '__main__':
     main()

phonopy_vibspec/spectra.py

@@ -0,0 +1,149 @@
+import pathlib
+from typing import List, Optional
+
+import h5py
+import numpy
+from numpy.typing import NDArray
+
+
+class RamanSpectrum:
+    def __init__(
+        self,
+        # input
+        modes: List[int],
+        frequencies: List[float],
+        irrep_labels: Optional[List[str]] = None,
+        dalpha_dq: Optional[NDArray[float]] = None,
+        # calculations
+        stencil: Optional[list] = None,
+        steps: List[float] = None,
+        dielectrics: Optional[NDArray] = None,
+    ):
+        assert irrep_labels is None or len(irrep_labels) == len(modes)
+        assert len(frequencies) == len(modes)
+        assert dalpha_dq is None or dalpha_dq.shape[0] == len(modes)
+
+        assert steps is None or len(steps) == len(modes)
+        assert dielectrics is None or dielectrics.shape[0] == len(modes)
+
+        self.modes = modes
+        self.frequencies = frequencies
+        self.irrep_labels = irrep_labels
+        self.dalpha_dq = dalpha_dq
+
+        self.steps = steps
+        self.dielectrics = dielectrics
+        self.stencil = stencil
+
+    def __len__(self):
+        return len(self.modes)
+
+    @classmethod
+    def from_hdf5(cls, path: pathlib.Path):
+        with h5py.File(path) as f:
+            if 'input' not in f:
+                raise Exception('missing `input` group')
+
+            modes = f['input/modes'][:]
+            frequencies = f['input/frequencies'][:]
+            irrep_labels = None
+            dalpha_dq = None
+            if 'input/irrep_labels' in f:
+                irrep_labels = f['input/irrep_labels'][:]
+            if 'input/dalpha_dq' in f:
+                dalpha_dq = f['input/dalpha_dq'][:]
+
+            stencil = None
+            steps = None
+            dielectrics = None
+
+            if 'numerical_differentiation' in f:
+                stencil = f['numerical_differentiation/stencil'][:]
+                steps = f['numerical_differentiation/steps'][:]
+
+                if 'numerical_differentiation/dielectrics' in f:
+                    dielectrics = f['numerical_differentiation/dielectrics'][:]
+
+            return cls(
+                modes=modes, frequencies=frequencies, irrep_labels=irrep_labels, dalpha_dq=dalpha_dq,
+                stencil=stencil, steps=steps, dielectrics=dielectrics
+            )
+
+    def to_hdf5(self, path: pathlib.Path):
+        with h5py.File(path, 'w') as f:
+            f.attrs['spectrum'] = 'Raman'
+            f.attrs['version'] = 1
+
+            g_input = f.create_group('input')
+            g_input.create_dataset('modes', data=self.modes)
+            g_input.create_dataset('frequencies', data=self.frequencies)
+
+            if self.irrep_labels:
+                g_input.create_dataset('irrep_labels', data=self.irrep_labels)
+
+            if self.dalpha_dq:
+                g_input.create_dataset('dalpha_dq', data=self.dalpha_dq)
+
+            g_calcs = f.create_group('numerical_differentiation')
+            g_calcs.create_dataset('stencil', data=numpy.array(self.stencil))
+
+            if self.steps:
+                g_calcs.create_dataset('steps', data=self.steps)
+
+            if self.dielectrics is not None:
+                g_calcs.create_dataset('dielectrics', data=self.dielectrics)
+
+
+class InfraredSpectrum:
+    def __init__(
+        self,
+        # input
+        modes: List[int],
+        frequencies: List[float],
+        dmu_dq: NDArray[float],
+        irrep_labels: Optional[List[str]] = None,
+    ):
+        assert irrep_labels is None or len(irrep_labels) == len(modes)
+        assert len(frequencies) == len(modes)
+        assert dmu_dq.shape[0] == len(modes)
+
+        self.modes = modes
+        self.frequencies = frequencies
+        self.irrep_labels = irrep_labels
+        self.dmu_dq = dmu_dq
+
+    def __len__(self):
+        return len(self.modes)
+
+    @classmethod
+    def from_hdf5(cls, path: pathlib.Path):
+        with h5py.File(path) as f:
+            assert f.attrs['spectrum'] == 'IR'
+
+            if f.attrs['version'] > 1:
+                raise Exception('unsupported version')
+
+            if 'input' not in f:
+                raise Exception('missing `input` group')
+
+            modes = f['input/modes'][:]
+            frequencies = f['input/frequencies'][:]
+            dmu_dq = f['input/dmu_dq'][:]
+            irrep_labels = None
+            if 'input/irrep_labels' in f:
+                irrep_labels = f['input/irrep_labels'][:]
+
+            return cls(modes=modes, frequencies=frequencies, irrep_labels=irrep_labels, dmu_dq=dmu_dq)
+
+    def to_hdf5(self, path: pathlib.Path):
+        with h5py.File(path, 'w') as f:
+            f.attrs['spectrum'] = 'IR'
+            f.attrs['version'] = 1
+
+            g_input = f.create_group('input')
+            g_input.create_dataset('modes', data=self.modes)
+            g_input.create_dataset('frequencies', data=self.frequencies)
+            g_input.create_dataset('dmu_dq', data=self.dmu_dq)
+
+    def compute_intensities(self) -> NDArray[float]:
+        return (self.dmu_dq ** 2).sum(axis=1)