implemented InterpolationModel.add_dof_from_pymatgen

ramannoodle/pmodel/interpolation.py

@@ -29,10 +29,19 @@
     verify_ndarray_shape,
     DOFWarning,
     UserError,
+    verify_list_len,
+    get_pymatgen_missing_error,
 )
 import ramannoodle.io.generic as generic_io
 from ramannoodle.io.utils import pathify_as_list
 
+PYMATGEN_PRESENT = True
+try:
+    import ramannoodle.io.pymatgen as pymatgen_io
+    import pymatgen.core
+except ImportError:
+    PYMATGEN_PRESENT = False
+
 
 def get_amplitude(
     cart_basis_vector: NDArray[np.float64],
@@ -516,6 +525,107 @@ def add_dof_from_files(
             interpolation_order,
         )
 
+    def _get_displacement_amplitudes_pymatgen(
+        self,
+        pymatgen_structures: "list[pymatgen.core.Structure]",
+    ) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+        """Get displacement and amplitudes from pymatgen Structures.
+
+        Parameters
+        ----------
+        pymatgen_structures
+            List of length M.
+
+        Returns
+        -------
+        :
+            0.  displacements -- (fractional) Array with shape (M,N,3) where N is the
+                number of atoms.
+            #.  amplitudes -- Array with shape (M,2).
+
+        Raises
+        ------
+        InvalidDOFException
+            DOF assembled from supplied Structures was invalid. See :meth:`add_dof` for
+            restrictions.
+        """
+        if not PYMATGEN_PRESENT:
+            raise get_pymatgen_missing_error()
+
+        # Calculate displacements and basis vector
+        displacements = []
+        for i, structure in enumerate(pymatgen_structures):
+            positions = pymatgen_io.get_positions(structure)
+            try:
+                displacement = calc_displacement(
+                    self._ref_structure.positions,
+                    positions,
+                )
+            except ValueError as exc:
+                invalid = f"pymatgen_structures[{i}]"
+                raise InvalidDOFException(f"incompatible structure: {invalid}") from exc
+            displacements.append(displacement)
+        result = is_collinear_with_all(displacements[0], displacements)
+        if result != -1:
+            invalid = f"pymatgen_structures[{result}]"
+            raise InvalidDOFException(f"displacement ({invalid}) is not collinear")
+        cart_basis_vector = self._ref_structure.get_cart_displacement(displacements[0])
+        cart_basis_vector /= np.linalg.norm(cart_basis_vector)
+
+        # Calculate amplitudes
+        amplitudes = []
+        for displacement in displacements:
+            cart_displacement = self._ref_structure.get_cart_displacement(displacement)
+            amplitudes.append(get_amplitude(cart_basis_vector, cart_displacement))
+        return (
+            np.array(displacements),
+            np.array(amplitudes),
+        )
+
+    def add_dof_from_pymatgen(
+        self,
+        pymatgen_structures: list[pymatgen.core.Structure],
+        polarizabilities: NDArray[np.float64],
+        interpolation_order: int,
+    ) -> None:
+        """
+        Add a degree of freedom from a list of pymatgen Structures and polarizabilities.
+
+        Parameters
+        ----------
+        pymatgen_structures
+            List of length M.
+        polarizabilities
+            Array with shape (M,3,3).
+        interpolation_order
+            Must be less than the number of total number of amplitudes after
+            symmetry considerations.
+
+        Raises
+        ------
+        InvalidDOFException
+            DOF assembled from supplied structures was invalid. See :meth:`add_dof` for
+            restrictions.
+        UserError
+            pymatgen is not installed.
+
+        """
+        if not PYMATGEN_PRESENT:
+            raise get_pymatgen_missing_error()
+        verify_list_len("pymatgen_structures", pymatgen_structures, None)
+
+        # Checks displacements
+        displacements, amplitudes = self._get_displacement_amplitudes_pymatgen(
+            pymatgen_structures
+        )
+        # Checks amplitudes
+        self.add_dof(
+            self.ref_structure.get_cart_displacement(displacements[0]),
+            amplitudes,
+            polarizabilities,
+            interpolation_order,
+        )
+
     def _read_dof(
         self, filepaths: str | Path | list[str] | list[Path], file_format: str
     ) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:

test/tests/test_phonon_spectrum.py

@@ -4,9 +4,12 @@
 from typing import Type
 import re
 
+import pytest
+
 import numpy as np
 from numpy.typing import NDArray
-import pytest
+
+import pymatgen.core
 
 import ramannoodle.io.generic
 from ramannoodle.pmodel.interpolation import InterpolationModel
@@ -105,6 +108,89 @@ def test_interpolation_spectrum(
         assert np.allclose(intensities, known_intensities, atol=1e-4)
 
 
+@pytest.mark.parametrize(
+    "outcar_ref_structure_fixture,data_directory,dof_eps_outcars",
+    [
+        (
+            "test/data/TiO2/phonons_OUTCAR",
+            "test/data/TiO2",
+            [
+                ["Ti5_0.1z_eps_OUTCAR", "Ti5_0.2z_eps_OUTCAR"],
+                ["Ti5_0.1x_eps_OUTCAR", "Ti5_0.2x_eps_OUTCAR"],
+                [
+                    "O43_0.1z_eps_OUTCAR",
+                    "O43_0.2z_eps_OUTCAR",
+                    "O43_m0.1z_eps_OUTCAR",
+                    "O43_m0.2z_eps_OUTCAR",
+                ],
+                ["O43_0.1x_eps_OUTCAR", "O43_0.2x_eps_OUTCAR"],
+                ["O43_0.1y_eps_OUTCAR", "O43_0.2y_eps_OUTCAR"],
+            ],
+        ),
+    ],
+    indirect=["outcar_ref_structure_fixture"],
+)
+def test_interpolation_spectrum_pymatgen(
+    outcar_ref_structure_fixture: ReferenceStructure,
+    data_directory: str,
+    dof_eps_outcars: list[str],
+) -> None:
+    """Test a full spectrum calculation using InterpolationModel."""
+    # Setup model
+    ref_structure = outcar_ref_structure_fixture
+    _, polarizability = ramannoodle.io.generic.read_positions_and_polarizability(
+        f"{data_directory}/ref_eps_OUTCAR", file_format="outcar"
+    )
+    model = InterpolationModel(ref_structure, polarizability)
+    for outcar_names in dof_eps_outcars:
+        pymatgen_structures = []
+        polarizabilities = []
+        for outcar_name in outcar_names:
+            positions, polarizability = (
+                ramannoodle.io.generic.read_positions_and_polarizability(
+                    f"{data_directory}/{outcar_name}", file_format="outcar"
+                )
+            )
+            ramannoodle.io.vasp.poscar.write_structure(
+                ref_structure.lattice,
+                ref_structure.atomic_numbers,
+                positions,
+                "test/data/scratch/PYMATGEN_INTERPOLATION_POSCAR",
+                overwrite=True,
+            )
+            pymatgen_structures.append(
+                pymatgen.core.Structure.from_file(
+                    "test/data/scratch/PYMATGEN_INTERPOLATION_POSCAR"
+                )
+            )
+            polarizabilities.append(polarizability)
+
+        model.add_dof_from_pymatgen(
+            pymatgen_structures, np.array(polarizabilities), interpolation_order=2
+        )
+
+    _validate_polarizabilities(model, data_directory)
+
+    # Spectrum test
+    with np.load(f"{data_directory}/known_spectrum.npz") as known_spectrum:
+        phonons = ramannoodle.io.generic.read_phonons(
+            f"{data_directory}/phonons_OUTCAR", file_format="outcar"
+        )
+        spectrum = phonons.get_raman_spectrum(model)
+        wavenumbers, intensities = spectrum.measure(
+            laser_correction=True,
+            laser_wavelength=532,
+            bose_einstein_correction=True,
+            temperature=300,
+        )
+
+        known_wavenumbers = known_spectrum["wavenumbers"]
+        known_intensities = known_spectrum["intensities"]
+
+        assert np.allclose(wavenumbers, known_wavenumbers)
+        assert np.allclose(intensities, known_intensities, atol=1e-4)
+
+
 @pytest.mark.parametrize(
     "outcar_ref_structure_fixture,data_directory,dof_eps_outcars",
     [