diff --git a/tests/test_io.py b/tests/test_io.py index 63616273..edbcf9f0 100644 --- a/tests/test_io.py +++ b/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), diff --git a/torch_sim/io.py b/torch_sim/io.py index ee5a46e5..6f3b1456 100644 --- a/torch_sim/io.py +++ b/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,6 +464,27 @@ 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, @@ -438,6 +492,8 @@ def structures_to_state( pbc=pbc_state, atomic_numbers=atomic_numbers, system_idx=system_idx, + _system_extras=_system_extras, + _atom_extras=_atom_extras, ) diff --git a/torch_sim/state.py b/torch_sim/state.py index 768e0f9b..402cb990 100644 --- a/torch_sim/state.py +++ b/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. diff --git a/torch_sim/trajectory.py b/torch_sim/trajectory.py index ef833a49..ec4a1141 100644 --- a/torch_sim/trajectory.py +++ b/torch_sim/trajectory.py @@ -1158,7 +1158,7 @@ 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 @@ -1166,6 +1166,7 @@ def get_structure(self, frame: int = -1) -> Any: 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.