From 414e5f92b9a479df5bc2fbe4666fb5b7689d6908 Mon Sep 17 00:00:00 2001 From: Johannes Spies <13813209+johannes-spies@users.noreply.github.com> Date: Wed, 15 Jul 2026 11:09:55 +0200 Subject: [PATCH 1/2] Fix FlashMD momenta and masses attachment --- src/metatrain/utils/data/dataset.py | 41 +++---------- tests/utils/data/test_dataset.py | 92 +++++++++++++++++++++++++++++ 2 files changed, 101 insertions(+), 32 deletions(-) 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..139168ee55 100644 --- a/tests/utils/data/test_dataset.py +++ b/tests/utils/data/test_dataset.py @@ -1201,3 +1201,95 @@ 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 test_memmap_momenta_attached(tmp_path): + """Momenta need to be attached and can be loaded.""" + target_options, _ = _write_minimal_memmap(tmp_path, ns=1) + # FlashMD stores the current momenta of every atom (1 atom per system here) + np.array([[1.0, 2.0, 3.0]], dtype="float32").tofile(tmp_path / "momenta.bin") + + 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 + ns = 2 + na = np.array([0, 2, 5], dtype=np.int64) # cumulative atom counts + np.save(tmp_path / "ns.npy", ns) + np.save(tmp_path / "na.npy", na) + np.zeros((5, 3), dtype="float32").tofile(tmp_path / "x.bin") + np.array([1, 1, 6, 6, 8], dtype="int32").tofile(tmp_path / "a.bin") + np.zeros((ns, 3, 3), dtype="float32").tofile(tmp_path / "c.bin") + # per-atom momenta; the x component encodes the global atom index 0..4 + momenta = np.zeros((5, 3), dtype="float32") + momenta[:, 0] = np.arange(5) + momenta.tofile(tmp_path / "momenta.bin") + # per-system energy target + np.array([1.0, 2.0], dtype="float32").reshape(ns, 1).tofile(tmp_path / "e.bin") + + target_options = { + "energy": { + "key": "e", + "sample_kind": "system", + "num_subtargets": 1, + "type": "scalar", + "quantity": "energy", + "forces": False, + "stress": False, + "virial": False, + } + } + + 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``. + """ + target_options, _ = _write_minimal_memmap(tmp_path, ns=2) + # FlashMD stores the mass of every atom (1 atom per system here) + np.array([1.0, 12.0], dtype="float32").tofile(tmp_path / "masses.bin") + + 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] From ce29bf2c52aca8885757ad217c1136992be5cbda Mon Sep 17 00:00:00 2001 From: Johannes Spies <13813209+johannes-spies@users.noreply.github.com> Date: Fri, 17 Jul 2026 15:20:28 +0200 Subject: [PATCH 2/2] Use MemmapWriter in test --- tests/utils/data/test_dataset.py | 133 ++++++++++++++++++++++++------- 1 file changed, 105 insertions(+), 28 deletions(-) diff --git a/tests/utils/data/test_dataset.py b/tests/utils/data/test_dataset.py index 139168ee55..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" @@ -1208,11 +1210,99 @@ def test_load_indices_file_not_found(tmp_path): # ============================================================ +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.""" - target_options, _ = _write_minimal_memmap(tmp_path, ns=1) # FlashMD stores the current momenta of every atom (1 atom per system here) - np.array([[1.0, 2.0, 3.0]], dtype="float32").tofile(tmp_path / "momenta.bin") + 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) @@ -1231,32 +1321,15 @@ def test_memmap_momenta_attached_in_batch(tmp_path): 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 - ns = 2 - na = np.array([0, 2, 5], dtype=np.int64) # cumulative atom counts - np.save(tmp_path / "ns.npy", ns) - np.save(tmp_path / "na.npy", na) - np.zeros((5, 3), dtype="float32").tofile(tmp_path / "x.bin") - np.array([1, 1, 6, 6, 8], dtype="int32").tofile(tmp_path / "a.bin") - np.zeros((ns, 3, 3), dtype="float32").tofile(tmp_path / "c.bin") # per-atom momenta; the x component encodes the global atom index 0..4 - momenta = np.zeros((5, 3), dtype="float32") + momenta = np.zeros((5, 3)) momenta[:, 0] = np.arange(5) - momenta.tofile(tmp_path / "momenta.bin") - # per-system energy target - np.array([1.0, 2.0], dtype="float32").reshape(ns, 1).tofile(tmp_path / "e.bin") - - target_options = { - "energy": { - "key": "e", - "sample_kind": "system", - "num_subtargets": 1, - "type": "scalar", - "quantity": "energy", - "forces": False, - "stress": False, - "virial": False, - } - } + 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())) @@ -1280,9 +1353,13 @@ 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``. """ - target_options, _ = _write_minimal_memmap(tmp_path, ns=2) # FlashMD stores the mass of every atom (1 atom per system here) - np.array([1.0, 12.0], dtype="float32").tofile(tmp_path / "masses.bin") + 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)