added ARTModel.add_dof_from_pymatgen

ramannoodle/pmodel/art.py

@@ -18,6 +18,12 @@
     UserError,
 )
 
+PYMATGEN_PRESENT = True
+try:
+    import pymatgen.core
+except (UserError, ImportError):
+    PYMATGEN_PRESENT = False
+
 
 def _get_directions_str(directions: list[NDArray[np.float64]]) -> str:
     """Return easy-to-read string for directions."""
@@ -109,6 +115,25 @@ def add_dof_from_files(
             "add_dof_from_files should not be used; use add_art_from_files instead"
         )
 
+    def add_dof_from_pymatgen(
+        self,
+        pymatgen_structures: "list[pymatgen.core.Structure]",
+        polarizabilities: NDArray[np.float64],
+        interpolation_order: int,
+    ) -> None:
+        """Disable add_dof_from_pymatgen.
+
+        Raises
+        ------
+        UserError
+
+        :meta private:
+        """
+        raise UserError(
+            "add_dof_from_pymatgen should not be used; use add_art_from_pymatgen "
+            "instead"
+        )
+
     def add_art(
         self,
         atom_index: int,
@@ -189,8 +214,6 @@ def add_art_from_files(
             Supports ``"outcar"`` and ``"vasprun.xml"``. If dummy model, supports
             ``"poscar"`` and ``"xdatcar"`` as well (see :ref:`Supported formats`).
 
-
-
         Raises
         ------
         FileNotFoundError
@@ -228,6 +251,58 @@ def add_art_from_files(
             basis_vectors_to_add, interpolation_xs, interpolation_ys, 1
         )
 
+    def add_art_from_pymatgen(
+        self,
+        pymatgen_structures: list[pymatgen.core.Structure],
+        polarizabilities: NDArray[np.float64],
+    ) -> None:
+        """
+        Add an atomic Raman tensor (ART) from pymatgen Structures and polarizabilities.
+
+        Required directions, amplitudes, and polarizabilities are automatically
+        determined from provided structures. Structures should be chosen such that the
+        resulting ARTs are valid under the same restrictions of :meth:`add_art`.
+
+        Parameters
+        ----------
+        pymatgen_structures
+            List of length M.
+        polarizabilities
+            Array with shape (M,3,3).
+
+        Raises
+        ------
+        InvalidDOFException
+            ART assembled from supplied Structures was invalid. See :meth:`add_art` for
+            restrictions.
+        """
+        displacements, amplitudes = super()._get_displacement_amplitudes_pymatgen(
+            pymatgen_structures
+        )
+
+        # Check whether only one atom is displaced.
+        _displacement = displacements[0].copy()
+        atom_index = int(np.argmax(np.sum(_displacement**2, axis=1)))
+        _displacement[atom_index] = np.zeros(3)
+        if not np.allclose(_displacement, 0.0, atol=1e-6):
+            raise InvalidDOFException("multiple atoms displaced simultaneously")
+
+        # Checks displacement
+        basis_vectors_to_add, interpolation_xs, interpolation_ys = super()._get_dof(
+            displacements[0], amplitudes, polarizabilities, False
+        )
+
+        num_amplitudes = len(interpolation_xs[0])
+        if num_amplitudes != 2:
+            raise InvalidDOFException(
+                f"wrong number of amplitudes: {num_amplitudes} != 2"
+            )
+
+        # Checks amplitudes
+        super()._construct_and_add_interpolations(
+            basis_vectors_to_add, interpolation_xs, interpolation_ys, 1
+        )
+
     def get_specification_tuples(
         self,
     ) -> list[tuple[int, list[int], list[NDArray[np.float64]]]]:

test/tests/test_phonon_spectrum.py

@@ -248,6 +248,83 @@ def test_art_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.1x_eps_OUTCAR"],
+                [
+                    "O43_0.1z_eps_OUTCAR",
+                    "O43_m0.1z_eps_OUTCAR",
+                ],
+                ["O43_0.1x_eps_OUTCAR"],
+                ["O43_0.1y_eps_OUTCAR"],
+            ],
+        ),
+    ],
+    indirect=["outcar_ref_structure_fixture"],
+)
+def test_art_spectrum_pymatgen(
+    outcar_ref_structure_fixture: ReferenceStructure,
+    data_directory: str,
+    dof_eps_outcars: list[str],
+) -> None:
+    """Test a full spectrum calculation using ARTModel and pymatgen."""
+    # 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 = ARTModel(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_ART_POSCAR",
+                overwrite=True,
+            )
+            pymatgen_structures.append(
+                pymatgen.core.Structure.from_file(
+                    "test/data/scratch/PYMATGEN_ART_POSCAR"
+                )
+            )
+            polarizabilities.append(polarizability)
+
+        model.add_art_from_pymatgen(pymatgen_structures, np.array(polarizabilities))
+
+    # Spectrum test
+    with np.load(f"{data_directory}/known_art_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,atoms_to_mask,"
     "known_spectrum_file",