diff --git a/src/metatrain/utils/data/dataset.py b/src/metatrain/utils/data/dataset.py index d3785245ce..3101ca0da6 100644 --- a/src/metatrain/utils/data/dataset.py +++ b/src/metatrain/utils/data/dataset.py @@ -1249,12 +1249,15 @@ def __getitem__(self, i: int) -> Any: samples=Labels( names=["system", "atom"], values=torch.tensor( - [[i, j] for j in range(self.na[i], self.na[i + 1])], + [ + [i, j] + for j in range(self.na[i + 1] - self.na[i]) + ], dtype=torch.int32, ), ), components=[Labels.range("xyz", 3)], - properties=Labels.single(), + properties=Labels.range("momentum", 1), ) ], ), @@ -1270,7 +1273,10 @@ def __getitem__(self, i: int) -> Any: samples=Labels( names=["system", "atom"], values=torch.tensor( - [[i, j] for j in range(self.na[i], self.na[i + 1])], + [ + [i, j] + for j in range(self.na[i + 1] - self.na[i]) + ], dtype=torch.int32, ), ), @@ -1384,35 +1390,6 @@ def __getitem__(self, i: int) -> Any: ) target_dict[target_key] = target_tensormap - momenta = getattr(self, "momenta", None) - if momenta is not None: - momenta = torch.tensor( - momenta[self.na[i] : self.na[i + 1]], dtype=torch.float64 - ) - system.add_data( - "momentum", - TensorMap( - keys=Labels.single(), - blocks=[ - TensorBlock( - values=momenta.unsqueeze(-1), - samples=Labels( - names=["system", "atom"], - values=torch.tensor( - [ - [i, j] - for j in range(self.na[i + 1] - self.na[i]) - ], - dtype=torch.int32, - ), - ), - components=[Labels.range("xyz", 3)], - properties=Labels.range("momentum", 1), - ), - ], - ), - ) - # Build extra_data TensorMaps returned in the sample and forwarded to # the `extra` argument of CollateFn callables extra_data_dict = {} diff --git a/tests/utils/data/test_dataset.py b/tests/utils/data/test_dataset.py index e84607b317..a436e8af86 100644 --- a/tests/utils/data/test_dataset.py +++ b/tests/utils/data/test_dataset.py @@ -4,6 +4,7 @@ import pytest import torch from metatensor.torch import Labels, TensorBlock, TensorMap +from metatomic.torch import System from omegaconf import OmegaConf from metatrain.utils.data import ( @@ -22,6 +23,7 @@ unpack_batch, ) from metatrain.utils.data.dataset import MemmapDataset +from metatrain.utils.data.writers import MemmapWriter RESOURCES_PATH = Path(__file__).parents[2] / "resources" @@ -1201,3 +1203,170 @@ def test_load_indices_file_not_found(tmp_path): """Missing file raises ValueError.""" with pytest.raises(ValueError, match="not found"): load_indices("nonexistent.txt") + + +# ============================================================ +# MemmapDataset FlashMD momenta and masses tests +# ============================================================ + + +def _write_memmap_via_writer( + tmp_path, atoms_per_system, energy_values, momenta=None, masses=None +): + """Write a MemmapDataset directory using the ``MemmapWriter``. + + :param atoms_per_system: number of atoms per system. + :param energy_values: one energy value per system. + :param momenta: optional (n_atoms, 3) array/list of per-atom momenta. + :param masses: optional (n_atoms,) array/list of per-atom masses. + """ + systems = [ + System( + types=torch.ones(n, dtype=torch.int32), + positions=torch.zeros(n, 3), + cell=torch.zeros(3, 3), + pbc=torch.zeros(3, dtype=torch.bool), + ) + for n in atoms_per_system + ] + + n_systems = len(systems) + predictions = { + "e": TensorMap( + Labels.single(), + [ + TensorBlock( + values=torch.tensor(energy_values, dtype=torch.float64).reshape( + n_systems, 1 + ), + samples=Labels("system", torch.arange(n_systems).reshape(-1, 1)), + components=[], + properties=Labels.range("energy", 1), + ) + ], + ) + } + + system_atom_labels = torch.tensor( + [[i, j] for i, n in enumerate(atoms_per_system) for j in range(n)] + ) + if momenta is not None: + predictions["momenta"] = TensorMap( + Labels.single(), + [ + TensorBlock( + values=torch.tensor(momenta, dtype=torch.float64).unsqueeze(-1), + samples=Labels(["system", "atom"], system_atom_labels), + components=[Labels.range("xyz", 3)], + properties=Labels.range("momentum", 1), + ) + ], + ) + if masses is not None: + predictions["masses"] = TensorMap( + Labels.single(), + [ + TensorBlock( + values=torch.tensor(masses, dtype=torch.float64).reshape(-1, 1), + samples=Labels(["system", "atom"], system_atom_labels), + components=[], + properties=Labels.range("mass", 1), + ) + ], + ) + + writer = MemmapWriter(tmp_path) + writer.write(systems, predictions) + writer.finish() + + target_options = { + "energy": { + "key": "e", + "sample_kind": "system", + "num_subtargets": 1, + "type": "scalar", + "quantity": "energy", + "forces": False, + "stress": False, + "virial": False, + } + } + return target_options + + +def test_memmap_momenta_attached(tmp_path): + """Momenta need to be attached and can be loaded.""" + # FlashMD stores the current momenta of every atom (1 atom per system here) + target_options = _write_memmap_via_writer( + tmp_path, + atoms_per_system=[1], + energy_values=[1.0], + momenta=[[1.0, 2.0, 3.0]], + ) + + dataset = MemmapDataset(tmp_path, target_options) + + sample = dataset[0] # used to raise ValueError + + assert sample.system.known_data() == ["momentum"], ( + f"momenta must be registered exactly once, got {sample.system.known_data()}" + ) + values = sample.system.get_data("momentum").block().values + assert values.squeeze(-1).tolist() == [[1.0, 2.0, 3.0]] + + +def test_memmap_momenta_attached_in_batch(tmp_path): + """ + Every system in a collated batch must carry its own momenta exactly once, + with local (0-based) atom labels and its own slice of ``momenta.bin``. + """ + # 2 systems: system 0 has 2 atoms, system 1 has 3 atoms + # per-atom momenta; the x component encodes the global atom index 0..4 + momenta = np.zeros((5, 3)) + momenta[:, 0] = np.arange(5) + target_options = _write_memmap_via_writer( + tmp_path, + atoms_per_system=[2, 3], + energy_values=[1.0, 2.0], + momenta=momenta, + ) + + dataset = MemmapDataset(tmp_path, target_options) + collate_fn = CollateFn(list(target_options.keys())) + + systems, _, _ = unpack_batch(collate_fn([dataset[0], dataset[1]])) + + assert len(systems) == 2 + # system 0 owns atoms 0-1, system 1 owns atoms 2-4; labels stay local + expected = [([0, 1], [0.0, 1.0]), ([0, 1, 2], [2.0, 3.0, 4.0])] + for system, (atom_labels, x_components) in zip(systems, expected, strict=True): + assert system.known_data() == ["momentum"], ( + f"momenta must be registered exactly once, got {system.known_data()}" + ) + block = system.get_data("momentum").block() + assert block.samples.values[:, 1].tolist() == atom_labels + assert block.values.squeeze(-1)[:, 0].tolist() == x_components + + +def test_memmap_masses_attached(tmp_path): + """ + Masses must be attached with local (0-based) atom labels and the system's own + slice of ``masses.bin``. + """ + # FlashMD stores the mass of every atom (1 atom per system here) + target_options = _write_memmap_via_writer( + tmp_path, + atoms_per_system=[1, 1], + energy_values=[1.0, 2.0], + masses=[1.0, 12.0], + ) + + dataset = MemmapDataset(tmp_path, target_options) + + system = dataset[1].system + + assert system.known_data() == ["mass"] + block = system.get_data("mass").block() + # the atom label is local (0), not the global offset (1) + assert block.samples.values.tolist() == [[1, 0]] + assert block.values.squeeze(-1).tolist() == [12.0]