Add opt-in hook to exempt fields from the fp-dtype downcast

CHANGELOG.md

@@ -63,6 +63,22 @@
   paths; per-batch transforms run on the consumer thread after `Batch.from_data_list`.
   Transform failures are wrapped in `RuntimeError` with `transform[<i>]`
   breadcrumb and `__cause__` preserved.
+- **Opt-in precision-preserving fields** — `AtomicData` gained a public
+  `precision_preserving_keys` class attribute (empty by default, so behaviour is
+  unchanged) listing floating-point fields exempt from
+  `check_fp_dtype_consistency`'s cast to the positions dtype. Set it globally
+  (`AtomicData.precision_preserving_keys = frozenset({"energy"})`) or on a
+  subclass to keep a high-precision label — e.g. so an extensive total energy
+  (~1e4–1e5 eV) is not silently downcast to float32. The exemption
+  also holds across `add_system_property` / `add_node_property`, and survives
+  batching: `Batch.from_data_list` collates with `torch.cat` and the storage
+  layer respects an existing tensor's dtype, so the preserved precision carries
+  through the batch and device moves unchanged.
+- **`check_fp_dtype_consistency` cast warning is now deduplicated** — the warning
+  fires only the first time each distinct `(field, source_dtype, target_dtype)`
+  cast is seen in a process, instead of on every `AtomicData` construction (which
+  flooded logs in training loops). Casting behaviour is unchanged; only the
+  warning frequency is reduced.
 
 ### Models
 

nvalchemi/data/atomic_data.py

@@ -53,6 +53,23 @@ def _tensor_serialization(tensor: torch.Tensor) -> list[float | int | list]:
     return tensor.detach().cpu().tolist()
 
 
+def _warn_first(seen: set, key: object, message: str, *, stacklevel: int = 2) -> None:
+    """Emit ``message`` as a ``UserWarning`` the first time ``key`` is seen.
+
+    ``key`` is recorded in ``seen`` so repeats stay silent. ``stacklevel`` is
+    relative to the caller (this frame is added on top).
+    """
+    if key not in seen:
+        seen.add(key)
+        warnings.warn(message, UserWarning, stacklevel=stacklevel + 1)
+
+
+# Process-global record of (field, src_dtype, tgt_dtype) casts already warned
+# about, so check_fp_dtype_consistency warns once per cast rather than per
+# construction. Module-level so it stays out of the AtomicData schema.
+_FP_CAST_WARNED: set[tuple[str, str, str]] = set()
+
+
 class AtomicNumberTable:
     """
     Atomic number table
@@ -366,6 +383,12 @@ class AtomicData(BaseModel, DataMixin):
         }
     )
 
+    # FP fields exempt from the positions-dtype cast below. Empty by default. Set
+    # it globally (``AtomicData.precision_preserving_keys = frozenset({"energy"})``)
+    # or on a subclass to keep a high-precision label, e.g. a fp64 total energy
+    # (~1e4-1e5 eV, which fp32 would quantize to ~1e-2 eV).
+    precision_preserving_keys: ClassVar[frozenset[str]] = frozenset()
+
     # Pydantic configuration
     model_config: ClassVar[ConfigDict] = ConfigDict(
         arbitrary_types_allowed=True, validate_assignment=True, extra="allow"
@@ -462,31 +485,28 @@ def check_edge_consistency(self) -> AtomicData:
     @model_validator(mode="after")
     def check_fp_dtype_consistency(self) -> AtomicData:
         """
-        Ensures all floating point tensors are at the same precision
-        as the positions tensor.
+        Cast floating point tensors to the positions dtype. Fields in
+        ``precision_preserving_keys`` are exempt. Casting is unconditional, but
+        each distinct ``(field, src, tgt)`` cast warns only once per process.
         """
         dtype = self.positions.dtype
-        casted: list[str] = []
         for key in self.model_dump().keys():
+            if key in self.precision_preserving_keys:
+                continue
             value = getattr(self, key)
-            if isinstance(value, torch.Tensor):
-                tensor_dtype = value.dtype
-                if tensor_dtype.is_floating_point and tensor_dtype != dtype:
-                    # using __dict__ to avoid re-validation
-                    self.__dict__[key] = value.to(dtype)
-                    casted.append(key)
-        if casted:
-            casted.sort()
-            # Keep the warning attributed to the user's AtomicData(...) call
-            # instead of Pydantic's internal validation frames. This may need
-            # adjustment if Pydantic's construction stack changes.
-            warnings.warn(
-                f"AtomicData fields {casted} were cast from their original "
-                f"dtypes to {dtype} to match positions. "
-                f"Pass tensors with matching dtypes to silence this warning.",
-                UserWarning,
-                stacklevel=3,
-            )
+            if not (isinstance(value, torch.Tensor) and value.dtype.is_floating_point):
+                continue
+            if value.dtype != dtype:
+                # stacklevel=3 keeps the warning on the user's AtomicData(...) call
+                _warn_first(
+                    _FP_CAST_WARNED,
+                    (key, str(value.dtype), str(dtype)),
+                    f"AtomicData field '{key}' was cast from {value.dtype} to "
+                    f"{dtype} to match positions; pass a matching dtype to silence "
+                    f"this (warned once per field/dtype).",
+                    stacklevel=3,
+                )
+                self.__dict__[key] = value.to(dtype)  # __dict__ skips re-validation
         return self
 
     @model_validator(mode="after")

test/conftest.py

@@ -17,6 +17,19 @@
 import pytest
 import torch
 
+from nvalchemi.data import atomic_data
+
+
+@pytest.fixture(autouse=True)
+def _reset_fp_cast_warned() -> None:
+    """Clear the process-global fp-cast warning cache before each test.
+
+    ``check_fp_dtype_consistency`` dedups its cast warning per ``(field, dtype)``
+    across all constructions, so without this reset warning assertions would
+    depend on test order (an earlier test could consume the only warning for a
+    given field/dtype)."""
+    atomic_data._FP_CAST_WARNED.clear()
+
 
 @pytest.fixture(params=["cpu", "cuda"])
 def device(request) -> str:

test/data/test_atomic_data.py

@@ -692,6 +692,96 @@ def test_no_warning_when_dtypes_match(self):
         user_warnings = [w for w in caught if issubclass(w.category, UserWarning)]
         assert len(user_warnings) == 0
 
+    def test_energy_downcast_by_default(self):
+        """Default exempt set is empty: a fp64 energy is cast to positions dtype."""
+        e64 = torch.tensor([[-93873.600167206]], dtype=torch.float64)
+        data = AtomicData(
+            positions=torch.randn(2, 3, dtype=torch.float32),
+            atomic_numbers=torch.ones(2, dtype=torch.long),
+            energy=e64,
+        )
+        assert data.energy.dtype == torch.float32
+
+    def test_precision_preserving_keys_keeps_energy(self, monkeypatch):
+        """Listing a field in the public precision_preserving_keys keeps its
+        precision, including across validate_assignment re-runs on setattr."""
+        monkeypatch.setattr(
+            AtomicData, "precision_preserving_keys", frozenset({"energy"})
+        )
+        e64 = torch.tensor([[-41512.590527111]], dtype=torch.float64)
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            data = AtomicData(
+                positions=torch.randn(3, 3, dtype=torch.float32),
+                atomic_numbers=torch.ones(3, dtype=torch.long),
+                energy=e64,
+            )
+        assert data.energy.dtype == torch.float64
+        assert data.energy.item() == e64.item()  # bit-exact, no fp32 round-trip
+        msgs = [str(w.message) for w in caught if issubclass(w.category, UserWarning)]
+        assert all("energy" not in m for m in msgs)
+        data.add_system_property("dataset_id", torch.tensor([0], dtype=torch.long))
+        assert data.energy.dtype == torch.float64
+        assert data.energy.item() == e64.item()
+
+    def test_non_exempt_fp_field_still_cast(self, monkeypatch):
+        """Fields not in the exempt set are still cast to the positions dtype."""
+        monkeypatch.setattr(
+            AtomicData, "precision_preserving_keys", frozenset({"energy"})
+        )
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            data = AtomicData(
+                positions=torch.randn(2, 3, dtype=torch.float32),
+                atomic_numbers=torch.ones(2, dtype=torch.long),
+                forces=torch.randn(2, 3, dtype=torch.float64),
+            )
+        assert data.forces.dtype == torch.float32
+        msgs = [str(w.message) for w in caught if issubclass(w.category, UserWarning)]
+        assert any("forces" in m for m in msgs)
+
+    def test_cast_warning_fires_once_per_field_dtype(self):
+        """The cast warning is emitted once per (field, dtype); repeats are silent
+        even though every instance is still cast."""
+
+        def make():
+            return AtomicData(
+                positions=torch.randn(2, 3, dtype=torch.float32),
+                atomic_numbers=torch.ones(2, dtype=torch.long),
+                forces=torch.randn(2, 3, dtype=torch.float64),
+            )
+
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            first = make()
+            second = make()  # same (forces, float64 -> float32): no second warning
+            make()
+        # casting still happens for every instance, only the warning is deduped
+        assert first.forces.dtype == torch.float32
+        assert second.forces.dtype == torch.float32
+        user_warnings = [w for w in caught if issubclass(w.category, UserWarning)]
+        assert len(user_warnings) == 1
+        assert "forces" in str(user_warnings[0].message)
+
+    def test_distinct_field_dtype_casts_each_warn_once(self):
+        """A different field or a different source dtype warns on its own."""
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            # (forces, float64 -> float32)
+            AtomicData(
+                positions=torch.randn(2, 3, dtype=torch.float32),
+                atomic_numbers=torch.ones(2, dtype=torch.long),
+                forces=torch.randn(2, 3, dtype=torch.float64),
+            )
+            # distinct field (velocities) -> a separate first-time warning
+            AtomicData(
+                positions=torch.randn(2, 3, dtype=torch.float32),
+                atomic_numbers=torch.ones(2, dtype=torch.long),
+                velocities=torch.randn(2, 3, dtype=torch.float64),
+            )
+        user_warnings = [w for w in caught if issubclass(w.category, UserWarning)]
+        assert len(user_warnings) == 2
+
 
 # -----------------------------------------------------------------------------
 # from_atoms: cell and pbc handling

test/data/test_batch.py

@@ -1317,3 +1317,69 @@ def test_field_levels_fallback_still_system(self) -> None:
         # field_levels is provided but doesn't mention unknown_scalar
         batch = Batch.from_raw_dicts([d0, d1], field_levels={"some_other_key": "atom"})
         assert "unknown_scalar" in batch.keys["system"]
+
+
+# -----------------------------------------------------------------------------
+# Precision-preserving fields survive batching
+# -----------------------------------------------------------------------------
+class TestPrecisionPreservingBatch:
+    """A field opted out of the fp-dtype downcast at the AtomicData level keeps
+    its precision through ``Batch.from_data_list`` and device moves.
+
+    ``Batch`` carries whatever dtype the source ``AtomicData`` produced (collation
+    is ``torch.cat``, and the storage layer respects an existing tensor's dtype),
+    so the ``precision_preserving_keys`` exemption is honoured end-to-end.
+    """
+
+    def test_fp64_energy_survives_from_data_list(self, monkeypatch):
+        monkeypatch.setattr(
+            AtomicData, "precision_preserving_keys", frozenset({"energy"})
+        )
+        e64 = torch.tensor([[-93873.600167206]], dtype=torch.float64)
+        data_list = [
+            AtomicData(
+                positions=torch.randn(2, 3, dtype=torch.float32),
+                atomic_numbers=torch.ones(2, dtype=torch.long),
+                energy=e64.clone(),
+            )
+            for _ in range(3)
+        ]
+        batch = Batch.from_data_list(data_list)
+        assert batch.energy.dtype == torch.float64
+        # bit-exact: no fp32 round-trip anywhere in the collation path
+        assert batch.energy[0].item() == e64.item()
+
+    def test_fp64_energy_survives_device_move(self, monkeypatch):
+        monkeypatch.setattr(
+            AtomicData, "precision_preserving_keys", frozenset({"energy"})
+        )
+        e64 = torch.tensor([[-41512.590527111]], dtype=torch.float64)
+        batch = Batch.from_data_list(
+            [
+                AtomicData(
+                    positions=torch.randn(2, 3, dtype=torch.float32),
+                    atomic_numbers=torch.ones(2, dtype=torch.long),
+                    energy=e64.clone(),
+                )
+                for _ in range(2)
+            ]
+        )
+        moved = batch.to("cpu")
+        assert moved.energy.dtype == torch.float64
+        assert moved.energy[0].item() == e64.item()
+
+    def test_default_energy_batches_at_positions_dtype(self):
+        """Without the exemption, energy is downcast at the AtomicData level and
+        the batch faithfully carries that float32 (no schema-driven upcast)."""
+        e64 = torch.tensor([[-93873.600167206]], dtype=torch.float64)
+        with pytest.warns(UserWarning, match="energy"):
+            data_list = [
+                AtomicData(
+                    positions=torch.randn(2, 3, dtype=torch.float32),
+                    atomic_numbers=torch.ones(2, dtype=torch.long),
+                    energy=e64.clone(),
+                )
+                for _ in range(2)
+            ]
+        batch = Batch.from_data_list(data_list)
+        assert batch.energy.dtype == torch.float32