Enhance pymatgen io functions to support system and atom extras

tests/test_io.py

@@ -7,6 +7,7 @@
 import torch
 from ase import Atoms
 from ase.build import molecule
+from pymatgen.core import Lattice, Structure
 
 import torch_sim as ts
 from tests.conftest import DEVICE, DTYPE
@@ -172,11 +173,106 @@ def test_state_round_trip(
         assert torch.allclose(sim_state.masses, round_trip_state.masses)
 
 
+@pytest.mark.parametrize(
+    ("system_extras_map", "atom_extras_map", "expected_sys", "expected_atom"),
+    [
+        pytest.param(None, None, {}, {}, id="no-extras-by-default"),
+        pytest.param(
+            {"charge": "charge", "spin": "spin"},
+            None,
+            {"charge": 3.0, "spin": 2.0},
+            {},
+            id="system-extras-identity-map",
+        ),
+        pytest.param(
+            {"total_charge": "charge"},
+            None,
+            {"total_charge": 3.0},
+            {},
+            id="system-extras-rename",
+        ),
+        pytest.param(
+            None,
+            {"site_tags": "my_tags"},
+            {},
+            {"site_tags": [1.0, 2.0]},
+            id="atom-extras-rename",
+        ),
+    ],
+)
+def test_extras_map_import_pymatgen(
+    system_extras_map: dict[SystemExtras, str] | None,
+    atom_extras_map: dict[AtomExtras, str] | None,
+    expected_sys: dict[str, float],
+    expected_atom: dict[str, list[float]],
+) -> None:
+    """test how system_extras_map and atom_extras_map control which keys are
+    read and how they are renamed on import from pymatgen Structures.
+    """
+    struct = Structure(Lattice.cubic(3.0), ["Si", "Si"], [[0, 0, 0], [0.5, 0.5, 0.5]])
+    struct.properties["charge"] = 3.0
+    struct.properties["spin"] = 2.0
+    struct.add_site_property("my_tags", [1.0, 2.0])
+    state = ts.io.structures_to_state(
+        [struct],
+        DEVICE,
+        DTYPE,
+        system_extras_map=system_extras_map,
+        atom_extras_map=atom_extras_map,
+    )
+    if not expected_sys and not expected_atom:
+        assert not state.system_extras
+        assert not state.atom_extras
+    for key, val in expected_sys.items():
+        assert getattr(state, key)[0].item() == val
+    for key, vals in expected_atom.items():
+        assert getattr(state, key).shape == (len(vals),)
+
+
+def test_extras_map_missing_key_skipped_pymatgen() -> None:
+    """Missing pymatgen keys are silently skipped rather than defaulting to zero."""
+    struct = Structure(Lattice.cubic(3.0), ["Si"], [[0, 0, 0]])
+    state = ts.io.structures_to_state(
+        [struct], DEVICE, DTYPE, system_extras_map={"charge": "charge"}
+    )
+    assert not state.system_extras
+
+
+def test_extras_map_multi_system_pymatgen() -> None:
+    """System extras work across multiple structures with correct per-system values."""
+    struct1 = Structure(Lattice.cubic(3.0), ["Si"], [[0, 0, 0]])
+    struct2 = Structure(Lattice.cubic(4.0), ["Fe"], [[0, 0, 0]])
+    struct1.properties["charge"] = 1.0
+    struct2.properties["charge"] = -1.0
+    state = ts.io.structures_to_state(
+        [struct1, struct2], DEVICE, DTYPE, system_extras_map={"charge": "charge"}
+    )
+    assert state.charge.shape == (2,)
+    assert state.charge[0].item() == 1.0
+    assert state.charge[1].item() == -1.0
+
+
+def test_extras_map_export_roundtrip_pymatgen() -> None:
+    """System and atom extras round-trip through state_to_structures with rename."""
+    struct = Structure(Lattice.cubic(3.0), ["Si", "Si"], [[0, 0, 0], [0.5, 0.5, 0.5]])
+    struct.properties["charge"] = 5.0
+    struct.add_site_property("my_tags", [1.0, 2.0])
+    sys_map = {"total_charge": "charge"}
+    atom_map = {"site_tags": "my_tags"}
+    state = ts.io.structures_to_state(
+        [struct], DEVICE, DTYPE, system_extras_map=sys_map, atom_extras_map=atom_map
+    )
+    structures = ts.io.state_to_structures(
+        state, system_extras_map=sys_map, atom_extras_map=atom_map
+    )
+    assert structures[0].properties["charge"] == 5.0
+    np.testing.assert_allclose(structures[0].site_properties["my_tags"], [1.0, 2.0])
+    structures_no_map = ts.io.state_to_structures(state)
+    assert "charge" not in structures_no_map[0].properties
+
+
 def test_structures_to_state_disordered() -> None:
     """structures_to_state rejects disordered (partial occupancy) structures."""
-    pytest.importorskip("pymatgen")
-    from pymatgen.core import Lattice, Structure
-
     # Site with partial occupancy (Cu/Au solid solution) -> disordered
     disordered = Structure(
         Lattice.cubic(3.6),

torch_sim/io.py

@@ -123,11 +123,22 @@ def state_to_atoms(
     description="pymatgen: Python Materials Genomics",
     path="pymatgen",
 )
-def state_to_structures(state: ts.SimState) -> list[Structure]:
+def state_to_structures(  # noqa: C901
+    state: ts.SimState,
+    *,
+    system_extras_map: dict[SystemExtras, str] | None = None,
+    atom_extras_map: dict[AtomExtras, str] | None = None,
+) -> list[Structure]:
     """Convert a SimState to a list of Pymatgen Structure objects.
 
     Args:
         state (SimState): Batched state containing positions, cell, and atomic numbers
+        system_extras_map: Map of ``{ts_key: pymatgen_key}`` controlling which
+            ``_system_extras`` entries are written to ``structure.properties``.
+            ``None`` (default) means no extras are written.
+        atom_extras_map: Map of ``{ts_key: pymatgen_key}`` controlling which
+            ``_atom_extras`` entries are written to ``structure.site_properties``.
+            ``None`` (default) means no extras are written.
 
     Returns:
         list[Structure]: Pymatgen Structure objects, one per system
@@ -184,6 +195,19 @@ def state_to_structures(state: ts.SimState) -> list[Structure]:
             coords=system_positions,
             coords_are_cartesian=True,
         )
+
+        if system_extras_map:
+            for ts_key, pmg_key in system_extras_map.items():
+                if ts_key in state.system_extras:
+                    val = state.system_extras[ts_key][uniq_sys_idx].detach().cpu().numpy()
+                    struct.properties[pmg_key] = val
+
+        if atom_extras_map:
+            for ts_key, pmg_key in atom_extras_map.items():
+                if ts_key in state.atom_extras:
+                    val = state.atom_extras[ts_key][mask].detach().cpu().numpy()
+                    struct.add_site_property(pmg_key, val)
+
         structures.append(struct)
 
     return structures
@@ -360,10 +384,13 @@ def atoms_to_state(
     description="pymatgen: Python Materials Genomics",
     path="pymatgen",
 )
-def structures_to_state(
+def structures_to_state(  # noqa: C901
     structure: Structure | list[Structure],
     device: torch.device | None = None,
     dtype: torch.dtype | None = None,
+    *,
+    system_extras_map: dict[SystemExtras, str] | None = None,
+    atom_extras_map: dict[AtomExtras, str] | None = None,
 ) -> ts.SimState:
     """Create a SimState from pymatgen Structure(s).
 
@@ -373,6 +400,12 @@ def structures_to_state(
         device (torch.device): Device to create tensors on
         dtype (torch.dtype): Data type for tensors (typically torch.float32 or
             torch.float64)
+        system_extras_map: Map of ``{ts_key: pymatgen_key}`` controlling which
+            ``pymatgen.properties`` entries are read into ``_system_extras``.
+            ``None`` (default) means no extras are read.
+        atom_extras_map: Map of ``{ts_key: pymatgen_key}`` controlling which
+            ``pymatgen.site_properties`` entries are read into ``_atom_extras``.
+            ``None`` (default) means no extras are read.
 
     Returns:
         SimState: TorchSim SimState object.
@@ -431,13 +464,36 @@ def structures_to_state(
     pbc_state: torch.Tensor | list[bool] | bool = (
         list(pbc_struct) if isinstance(pbc_struct, (list, tuple)) else pbc_struct
     )
+
+    _system_extras: dict[str, torch.Tensor] = {}
+    if system_extras_map:
+        for ts_key, pmg_key in system_extras_map.items():
+            vals = [at.properties.get(pmg_key) for at in struct_list]
+            non_none = [v for v in vals if v is not None]
+            if len(non_none) == len(vals):
+                _system_extras[ts_key] = torch.tensor(
+                    np.array(non_none), dtype=dtype, device=device
+                )
+
+    _atom_extras: dict[str, torch.Tensor] = {}
+    if atom_extras_map:
+        for ts_key, pmg_key in atom_extras_map.items():
+            arrays = [at.site_properties.get(pmg_key) for at in struct_list]
+            non_none = [a for a in arrays if a is not None]
+            if len(non_none) == len(arrays):
+                _atom_extras[ts_key] = torch.tensor(
+                    np.concatenate(non_none), dtype=dtype, device=device
+                )
+
     return ts.SimState(
         positions=positions,
         masses=masses,
         cell=cell,
         pbc=pbc_state,
         atomic_numbers=atomic_numbers,
         system_idx=system_idx,
+        _system_extras=_system_extras,
+        _atom_extras=_atom_extras,
     )
 
 

torch_sim/state.py

@@ -671,13 +671,24 @@ def to_atoms(
             self, system_extras_map=system_extras_map, atom_extras_map=atom_extras_map
         )
 
-    def to_structures(self) -> list["Structure"]:
+    def to_structures(
+        self,
+        *,
+        system_extras_map: dict[SystemExtras, str] | None = None,
+        atom_extras_map: dict[AtomExtras, str] | None = None,
+    ) -> list["Structure"]:
         """Convert the SimState to a list of pymatgen Structure objects.
 
+        Args:
+            system_extras_map: Map of ``{ts_key: pymatgen_key}`` for system extras.
+            atom_extras_map: Map of ``{ts_key: pymatgen_key}`` for atom extras.
+
         Returns:
             list[Structure]: A list of pymatgen Structure objects, one per system
         """
-        return ts.io.state_to_structures(self)
+        return ts.io.state_to_structures(
+            self, system_extras_map=system_extras_map, atom_extras_map=atom_extras_map
+        )
 
     def to_phonopy(self) -> list["PhonopyAtoms"]:
         """Convert the SimState to a list of PhonopyAtoms objects.

torch_sim/trajectory.py

@@ -1158,14 +1158,15 @@ def return_prop(self: Self, prop: str, frame: int) -> np.ndarray:
 
         return arrays
 
-    def get_structure(self, frame: int = -1) -> Any:
+    def get_structure(self, frame: int = -1, **kwargs: Any) -> Any:
         """Get a pymatgen Structure object for a given frame.
 
         Converts the state at the specified frame to a pymatgen Structure object
         for analysis and visualization.
 
         Args:
             frame (int, optional): Frame index to retrieve. Defaults to -1 for last frame.
+            **kwargs: Additional keyword arguments to pass to ``state_to_structures``.
 
         Returns:
             Structure: Pymatgen Structure object for the specified frame
@@ -1175,7 +1176,9 @@ def get_structure(self, frame: int = -1) -> Any:
         """
         from torch_sim.io import state_to_structures
 
-        return state_to_structures(self.get_state(frame, device=torch.device("cpu")))[0]
+        return state_to_structures(
+            self.get_state(frame, device=torch.device("cpu")), **kwargs
+        )[0]
 
     def get_atoms(self, frame: int = -1, **kwargs: Any) -> "Atoms":
         """Get an ASE Atoms object for a given frame.