displaced structure writing

.gitignore

@@ -267,3 +267,4 @@ $RECYCLE.BIN/
 session
 test/Testing.ipynb
 test/data/temp
+test/data/temp.vasp

ramannoodle/io/vasp/poscar.py

@@ -229,7 +229,7 @@ def write_structure(  # pylint: disable=too-many-arguments
     open_mode = "w" if overwrite else "x"
     filepath = pathify(filepath)
 
-    label_str = repr(label) + "\n"  # Raw string
+    label_str = repr(label)[1:-1] + "\n"  # Raw string with quotes removed
     lattice_str = _get_lattice_str(lattice)
     symbols_str = _get_symbols_str(atomic_numbers)
     positions_str = _get_positions_str(positions)

ramannoodle/polarizability/art.py

@@ -141,9 +141,7 @@ def add_art(
             Provided ART was invalid.
 
         """
-        direction = self.ref_structure.get_fractional_displacement(
-            np.array([cart_direction])
-        )
+        direction = self.ref_structure.get_frac_displacement(np.array([cart_direction]))
         if not isinstance(atom_index, int):
             raise get_type_error("atom_index", atom_index, "int")
         try:

ramannoodle/polarizability/interpolation.py

@@ -400,9 +400,7 @@ def add_dof(  # pylint: disable=too-many-arguments
 
         """
         try:
-            displacement = self.ref_structure.get_fractional_displacement(
-                cart_displacement
-            )
+            displacement = self.ref_structure.get_frac_displacement(cart_displacement)
 
             parent_displacement = displacement / np.linalg.norm(displacement * 10.0)
         except TypeError as exc:

ramannoodle/structure/displace.py

@@ -0,0 +1,194 @@
+"""Routines for generating and writing displaced structure."""
+
+# Design note:
+# These routines are not implemented in ReferenceStructure to give
+# greater modularity. For example, when different displacement methods are added (such
+# as Monte Carlo rattling or random displacements), we'd rather not add to code to
+# Reference Structure. These functions stand alone, just like the IO functions.
+
+from pathlib import Path
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ramannoodle.structure.reference import ReferenceStructure
+from ramannoodle.structure.structure_utils import displace_positions
+from ramannoodle.exceptions import (
+    get_type_error,
+    get_shape_error,
+    verify_ndarray_shape,
+    verify_list_len,
+)
+from ramannoodle.io.io_utils import pathify_as_list
+import ramannoodle.io.generic as rn_io
+
+
+def get_displaced_positions(
+    ref_structure: ReferenceStructure,
+    cart_displacement: NDArray[np.float64],
+    amplitudes: NDArray[np.float64],
+) -> list[NDArray[np.float64]]:
+    """Return list of displaced positions given a displacement and amplitudes.
+
+    Parameters
+    ----------
+    ref_structure
+        reference structure of N atoms
+    cart_displacement
+        2D array with shape (N,3)
+    amplitudes
+        1D array with shape (M,)
+
+    Returns
+    -------
+    :
+        1D list of length M
+    """
+    try:
+        cart_displacement = cart_displacement / float(np.linalg.norm(cart_displacement))
+    except TypeError as err:
+        raise get_type_error("cart_displacement", cart_displacement, "ndarray") from err
+    verify_ndarray_shape("amplitudes", amplitudes, (None,))
+
+    positions = []
+    for amplitude in amplitudes:
+        try:
+            displacement = ref_structure.get_frac_displacement(
+                cart_displacement * amplitude
+            )
+        except ValueError as err:
+            raise get_shape_error(
+                "cart_displacement",
+                cart_displacement,
+                f"({len(ref_structure.positions)}, 3)",
+            ) from err
+        positions.append(displace_positions(ref_structure.positions, displacement))
+
+    return positions
+
+
+def write_displaced_structures(  # pylint: disable=too-many-arguments
+    ref_structure: ReferenceStructure,
+    cart_displacement: NDArray[np.float64],
+    amplitudes: NDArray[np.float64],
+    file_paths: str | Path | list[str] | list[Path],
+    file_format: str,
+    overwrite: bool = False,
+) -> None:
+    """Write displaced structures to files.
+
+    Parameters
+    ----------
+    ref_structure
+        reference structure of N atoms
+    cart_displacement
+        2D array with shape (N,3)
+    amplitudes
+        1D array with shape (M,)
+    file_paths
+    file_format
+        supports: "poscar"
+    overwrite
+    """
+    file_paths = pathify_as_list(file_paths)
+    position_list = get_displaced_positions(
+        ref_structure, cart_displacement, amplitudes
+    )
+    verify_list_len("file_paths", file_paths, len(position_list))
+
+    for position, filepath in zip(position_list, file_paths):
+        rn_io.write_structure(
+            ref_structure.lattice,
+            ref_structure.atomic_numbers,
+            position,
+            filepath,
+            file_format,
+            overwrite,
+        )
+
+
+def get_ast_displaced_positions(
+    ref_structure: ReferenceStructure,
+    atom_index: int,
+    cart_direction: NDArray[np.float64],
+    amplitudes: NDArray[np.float64],
+) -> list[NDArray[np.float64]]:
+    """Return list of displaced positions with an atom displaced along a direction.
+
+    Parameters
+    ----------
+    ref_structure
+        reference structure of N atoms
+    atom_index
+    cart_direction
+        1D array with shape (3,)
+    amplitudes
+        1D array with shape (M,)
+
+    Returns
+    -------
+    :
+        1D list of length M
+    """
+    try:
+        cart_direction = cart_direction / float(np.linalg.norm(cart_direction))
+    except TypeError as err:
+        raise get_type_error("cart_direction", cart_direction, "ndarray") from err
+    positions = []
+    for amplitude in amplitudes:
+        cart_displacement = ref_structure.positions * 0
+        try:
+            cart_displacement[atom_index] = cart_direction * amplitude
+        except IndexError as err:
+            raise IndexError(f"invalid atom_index: {atom_index}") from err
+        displacement = ref_structure.get_frac_displacement(cart_displacement)
+        positions.append(displace_positions(ref_structure.positions, displacement))
+
+    return positions
+
+
+def write_ast_displaced_structures(  # pylint: disable=too-many-arguments
+    ref_structure: ReferenceStructure,
+    atom_index: int,
+    cart_direction: NDArray[np.float64],
+    amplitudes: NDArray[np.float64],
+    file_paths: str | Path | list[str] | list[Path],
+    file_format: str,
+    overwrite: bool = False,
+) -> None:
+    """Return displaced structures with an atom displaced along a direction.
+
+    Parameters
+    ----------
+    ref_structure
+        reference structure of N atoms
+    atom_index
+    cart_direction
+        1D array with shape (3,)
+    amplitudes
+        1D array with shape (M,)
+    file_paths
+    file_format
+        supports: "poscar"
+    overwrite
+
+    Returns
+    -------
+    :
+        1D list of length M
+    """
+    file_paths = pathify_as_list(file_paths)
+    position_list = get_ast_displaced_positions(
+        ref_structure, atom_index, cart_direction, amplitudes
+    )
+    verify_list_len("file_paths", file_paths, len(position_list))
+
+    for position, filepath in zip(position_list, file_paths):
+        rn_io.write_structure(
+            ref_structure.lattice,
+            ref_structure.atomic_numbers,
+            position,
+            filepath,
+            file_format,
+            overwrite,
+        )

ramannoodle/structure/reference.py

@@ -268,7 +268,7 @@ def get_cart_displacement(
 
         return displacement @ self._lattice
 
-    def get_fractional_displacement(
+    def get_frac_displacement(
         self, cart_displacement: NDArray[np.float64]
     ) -> NDArray[np.float64]:
         """Convert a Cartesian displacement into fractional coordinates.

test/tests/test_displace.py

@@ -0,0 +1,57 @@
+"""Tests for structure displacement routines."""
+
+import numpy as np
+from numpy.typing import NDArray
+
+import pytest
+
+import ramannoodle.io.vasp as vasp_io
+from ramannoodle.structure.reference import ReferenceStructure
+from ramannoodle.structure.displace import write_ast_displaced_structures
+
+# pylint: disable=protected-access
+
+
+@pytest.mark.parametrize(
+    "outcar_ref_structure_fixture,atom_index,cart_direction,amplitude,outcar_known",
+    [
+        (
+            "test/data/TiO2/phonons_OUTCAR",
+            42,
+            np.array([0.0, 0, 1]),
+            0.1,
+            "test/data/TiO2/O43_0.1z_eps_OUTCAR",
+        ),
+        (
+            "test/data/TiO2/phonons_OUTCAR",
+            4,
+            np.array([10, 0, 0]),
+            -0.2,
+            "test/data/TiO2/Ti5_m0.2x_eps_OUTCAR",
+        ),
+    ],
+    indirect=["outcar_ref_structure_fixture"],
+)
+def test_write_ast_displaced_structures(
+    outcar_ref_structure_fixture: ReferenceStructure,
+    atom_index: int,
+    cart_direction: NDArray[np.float64],
+    amplitude: float,
+    outcar_known: str,
+) -> None:
+    """Test write_displaced_structures."""
+    amplitudes = np.array([amplitude])
+    write_ast_displaced_structures(
+        outcar_ref_structure_fixture,
+        atom_index,
+        cart_direction,
+        amplitudes,
+        ["test/data/temp.vasp"],
+        "poscar",
+        overwrite=True,
+    )
+
+    known_positions = vasp_io.outcar.read_positions(outcar_known)
+    assert np.isclose(
+        vasp_io.poscar.read_positions("test/data/temp.vasp"), known_positions
+    ).all()