Fix FlashMD momenta and masses attachment

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 = {}

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]