From 3e4f6b6c394a5173365c57523840d3273a89b52d Mon Sep 17 00:00:00 2001
From: Eric Qu <ericq@nvidia.com>
Date: Thu, 2 Jul 2026 14:30:18 -0700
Subject: [PATCH 1/3] data: add opt-in hook to exempt fields from the fp-dtype
 downcast

AtomicData.check_fp_dtype_consistency casts every floating-point tensor to
positions.dtype (typically float32) so the model computes in one precision.
That is the right default for compute inputs, but it also downcasts label
fields: e.g. total energy is extensive (~1e4-1e5 eV for large systems), so
float32 quantizes it to ~1e-2 eV and reconstructed energies come out
discrete/staircased downstream.

Add a _precision_preserving_keys ClassVar and skip those fields in the cast.
It is empty by default, so behavior is unchanged; a subclass can opt a
high-precision label out of the downcast, e.g.

    class MyData(AtomicData):
        _precision_preserving_keys = frozenset({"energy"})

The exemption holds across add_system_property / add_node_property, since
validate_assignment=True re-runs the validator on every setattr.

Signed-off-by: Eric Qu <ericq@nvidia.com>
---
 CHANGELOG.md                  |  7 +++++
 nvalchemi/data/atomic_data.py | 14 ++++++++--
 test/data/test_atomic_data.py | 49 +++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 07bc0682..e64f80c5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -48,6 +48,13 @@
   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
+  `_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. A subclass can set
+  it to keep a high-precision label — e.g. `frozenset({"energy"})` 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`.
 
 ### Models
 
diff --git a/nvalchemi/data/atomic_data.py b/nvalchemi/data/atomic_data.py
index 1bf66d62..c23ab4df 100644
--- a/nvalchemi/data/atomic_data.py
+++ b/nvalchemi/data/atomic_data.py
@@ -366,6 +366,12 @@ class AtomicData(BaseModel, DataMixin):
         }
     )
 
+    # FP fields exempt from the positions-dtype cast below. Empty by default (all
+    # fp tensors match positions); a subclass may set it to keep a high-precision
+    # label, e.g. {"energy"} so a fp64 total energy (~1e4-1e5 eV, which fp32 would
+    # quantize to ~1e-2 eV) is not silently downcast to the compute precision.
+    _precision_preserving_keys: ClassVar[frozenset[str]] = frozenset()
+
     # Pydantic configuration
     model_config: ClassVar[ConfigDict] = ConfigDict(
         arbitrary_types_allowed=True, validate_assignment=True, extra="allow"
@@ -462,12 +468,16 @@ 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 for single-precision
+        compute. Fields in ``_precision_preserving_keys`` are exempt (empty by
+        default); a subclass can opt a high-precision label (e.g. ``energy``) out
+        of the downcast.
         """
         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
diff --git a/test/data/test_atomic_data.py b/test/data/test_atomic_data.py
index 6c8dc10c..535c857f 100644
--- a/test/data/test_atomic_data.py
+++ b/test/data/test_atomic_data.py
@@ -20,6 +20,7 @@
 import textwrap
 import warnings
 from pathlib import Path
+from typing import ClassVar
 
 import numpy as np
 import pytest
@@ -626,6 +627,12 @@ def test_atomic_numbers_default_int32(self):
 # -----------------------------------------------------------------------------
 # dtype cast warning
 # -----------------------------------------------------------------------------
+class _KeepEnergyData(AtomicData):
+    """AtomicData subclass that opts ``energy`` out of the fp-dtype downcast."""
+
+    _precision_preserving_keys: ClassVar[frozenset[str]] = frozenset({"energy"})
+
+
 class TestDtypeCastWarning:
     """Tests for check_fp_dtype_consistency warning."""
 
@@ -692,6 +699,48 @@ 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_subclass_can_preserve_energy_precision(self):
+        """A subclass listing a field in _precision_preserving_keys keeps its
+        precision, including across validate_assignment re-runs on setattr."""
+        e64 = torch.tensor([[-41512.590527111]], dtype=torch.float64)
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            data = _KeepEnergyData(
+                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):
+        """Fields not in the exempt set are still cast to the positions dtype."""
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            data = _KeepEnergyData(
+                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)
+
 
 # -----------------------------------------------------------------------------
 # from_atoms: cell and pbc handling

From f24b0153d6cabf2f4bbae855007fe33e6b0c2944 Mon Sep 17 00:00:00 2001
From: Eric Qu <ericq@nvidia.com>
Date: Mon, 6 Jul 2026 18:43:44 -0700
Subject: [PATCH 2/3] Ensure batch works with fp64 flag; add first time warning
 for casting

---
 CHANGELOG.md                  | 10 +++++-
 nvalchemi/data/atomic_data.py | 54 ++++++++++++++++------------
 test/conftest.py              | 13 +++++++
 test/data/test_atomic_data.py | 42 ++++++++++++++++++++++
 test/data/test_batch.py       | 68 +++++++++++++++++++++++++++++++++++
 5 files changed, 163 insertions(+), 24 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e64f80c5..a8d8fc0c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -54,7 +54,15 @@
   `check_fp_dtype_consistency`'s cast to the positions dtype. A subclass can set
   it to keep a high-precision label — e.g. `frozenset({"energy"})` 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`.
+  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
 
diff --git a/nvalchemi/data/atomic_data.py b/nvalchemi/data/atomic_data.py
index c23ab4df..824110da 100644
--- a/nvalchemi/data/atomic_data.py
+++ b/nvalchemi/data/atomic_data.py
@@ -53,6 +53,17 @@ 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)
+
+
 class AtomicNumberTable:
     """
     Atomic number table
@@ -372,6 +383,10 @@ class AtomicData(BaseModel, DataMixin):
     # quantize to ~1e-2 eV) is not silently downcast to the compute precision.
     _precision_preserving_keys: ClassVar[frozenset[str]] = frozenset()
 
+    # Process-global (field, src_dtype, tgt_dtype) casts already warned about, so
+    # check_fp_dtype_consistency warns once per cast rather than per construction.
+    _fp_cast_warned: ClassVar[set[tuple[str, str, str]]] = set()
+
     # Pydantic configuration
     model_config: ClassVar[ConfigDict] = ConfigDict(
         arbitrary_types_allowed=True, validate_assignment=True, extra="allow"
@@ -468,35 +483,28 @@ def check_edge_consistency(self) -> AtomicData:
     @model_validator(mode="after")
     def check_fp_dtype_consistency(self) -> AtomicData:
         """
-        Cast floating point tensors to the positions dtype for single-precision
-        compute. Fields in ``_precision_preserving_keys`` are exempt (empty by
-        default); a subclass can opt a high-precision label (e.g. ``energy``) out
-        of the downcast.
+        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(
+                    self._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")
diff --git a/test/conftest.py b/test/conftest.py
index d3647882..906b304e 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -17,6 +17,19 @@
 import pytest
 import torch
 
+from nvalchemi.data.atomic_data import AtomicData
+
+
+@pytest.fixture(autouse=True)
+def _reset_fp_cast_warned() -> None:
+    """Clear the process-global fp-cast warning cache before each test.
+
+    ``AtomicData.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)."""
+    AtomicData._fp_cast_warned.clear()
+
 
 @pytest.fixture(params=["cpu", "cuda"])
 def device(request) -> str:
diff --git a/test/data/test_atomic_data.py b/test/data/test_atomic_data.py
index 535c857f..11db0e93 100644
--- a/test/data/test_atomic_data.py
+++ b/test/data/test_atomic_data.py
@@ -741,6 +741,48 @@ def test_non_exempt_fp_field_still_cast(self):
         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
diff --git a/test/data/test_batch.py b/test/data/test_batch.py
index 34b62e19..050d000f 100644
--- a/test/data/test_batch.py
+++ b/test/data/test_batch.py
@@ -16,6 +16,8 @@
 
 from __future__ import annotations
 
+from typing import ClassVar
+
 import pytest
 import torch
 
@@ -24,6 +26,12 @@
 from nvalchemi.data.level_storage import MultiLevelStorage, UniformLevelStorage
 
 
+class _KeepEnergyData(AtomicData):
+    """AtomicData subclass that opts ``energy`` out of the fp-dtype downcast."""
+
+    _precision_preserving_keys: ClassVar[frozenset[str]] = frozenset({"energy"})
+
+
 def _minimal_atomic_data(
     num_nodes: int = 4,
     num_edges: int = 0,
@@ -1317,3 +1325,63 @@ 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):
+        e64 = torch.tensor([[-93873.600167206]], dtype=torch.float64)
+        data_list = [
+            _KeepEnergyData(
+                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):
+        e64 = torch.tensor([[-41512.590527111]], dtype=torch.float64)
+        batch = Batch.from_data_list(
+            [
+                _KeepEnergyData(
+                    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

From 190ab5d93012dbbc7542fb8f4695f4b696be0492 Mon Sep 17 00:00:00 2001
From: Eric Qu <ericq@nvidia.com>
Date: Wed, 8 Jul 2026 06:58:13 -0700
Subject: [PATCH 3/3] Address comments

---
 CHANGELOG.md                  | 11 ++++++-----
 nvalchemi/data/atomic_data.py | 26 ++++++++++++++------------
 test/conftest.py              | 12 ++++++------
 test/data/test_atomic_data.py | 23 +++++++++++------------
 test/data/test_batch.py       | 24 +++++++++++-------------
 5 files changed, 48 insertions(+), 48 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 89c809f3..1c0a470e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -53,12 +53,13 @@
   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
-  `_precision_preserving_keys` class attribute (empty by default, so behaviour is
+- **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. A subclass can set
-  it to keep a high-precision label — e.g. `frozenset({"energy"})` so an extensive
-  total energy (~1e4–1e5 eV) is not silently downcast to float32. The exemption
+  `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
diff --git a/nvalchemi/data/atomic_data.py b/nvalchemi/data/atomic_data.py
index 824110da..36504e2b 100644
--- a/nvalchemi/data/atomic_data.py
+++ b/nvalchemi/data/atomic_data.py
@@ -64,6 +64,12 @@ def _warn_first(seen: set, key: object, message: str, *, stacklevel: int = 2) ->
         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
@@ -377,15 +383,11 @@ class AtomicData(BaseModel, DataMixin):
         }
     )
 
-    # FP fields exempt from the positions-dtype cast below. Empty by default (all
-    # fp tensors match positions); a subclass may set it to keep a high-precision
-    # label, e.g. {"energy"} so a fp64 total energy (~1e4-1e5 eV, which fp32 would
-    # quantize to ~1e-2 eV) is not silently downcast to the compute precision.
-    _precision_preserving_keys: ClassVar[frozenset[str]] = frozenset()
-
-    # Process-global (field, src_dtype, tgt_dtype) casts already warned about, so
-    # check_fp_dtype_consistency warns once per cast rather than per construction.
-    _fp_cast_warned: ClassVar[set[tuple[str, str, str]]] = set()
+    # 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(
@@ -484,12 +486,12 @@ def check_edge_consistency(self) -> AtomicData:
     def check_fp_dtype_consistency(self) -> AtomicData:
         """
         Cast floating point tensors to the positions dtype. Fields in
-        ``_precision_preserving_keys`` are exempt. Casting is unconditional, but
+        ``precision_preserving_keys`` are exempt. Casting is unconditional, but
         each distinct ``(field, src, tgt)`` cast warns only once per process.
         """
         dtype = self.positions.dtype
         for key in self.model_dump().keys():
-            if key in self._precision_preserving_keys:
+            if key in self.precision_preserving_keys:
                 continue
             value = getattr(self, key)
             if not (isinstance(value, torch.Tensor) and value.dtype.is_floating_point):
@@ -497,7 +499,7 @@ def check_fp_dtype_consistency(self) -> AtomicData:
             if value.dtype != dtype:
                 # stacklevel=3 keeps the warning on the user's AtomicData(...) call
                 _warn_first(
-                    self._fp_cast_warned,
+                    _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 "
diff --git a/test/conftest.py b/test/conftest.py
index 906b304e..f95f0d09 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -17,18 +17,18 @@
 import pytest
 import torch
 
-from nvalchemi.data.atomic_data import AtomicData
+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.
 
-    ``AtomicData.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)."""
-    AtomicData._fp_cast_warned.clear()
+    ``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"])
diff --git a/test/data/test_atomic_data.py b/test/data/test_atomic_data.py
index 11db0e93..b085efd0 100644
--- a/test/data/test_atomic_data.py
+++ b/test/data/test_atomic_data.py
@@ -20,7 +20,6 @@
 import textwrap
 import warnings
 from pathlib import Path
-from typing import ClassVar
 
 import numpy as np
 import pytest
@@ -627,12 +626,6 @@ def test_atomic_numbers_default_int32(self):
 # -----------------------------------------------------------------------------
 # dtype cast warning
 # -----------------------------------------------------------------------------
-class _KeepEnergyData(AtomicData):
-    """AtomicData subclass that opts ``energy`` out of the fp-dtype downcast."""
-
-    _precision_preserving_keys: ClassVar[frozenset[str]] = frozenset({"energy"})
-
-
 class TestDtypeCastWarning:
     """Tests for check_fp_dtype_consistency warning."""
 
@@ -709,13 +702,16 @@ def test_energy_downcast_by_default(self):
         )
         assert data.energy.dtype == torch.float32
 
-    def test_subclass_can_preserve_energy_precision(self):
-        """A subclass listing a field in _precision_preserving_keys keeps its
+    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 = _KeepEnergyData(
+            data = AtomicData(
                 positions=torch.randn(3, 3, dtype=torch.float32),
                 atomic_numbers=torch.ones(3, dtype=torch.long),
                 energy=e64,
@@ -728,11 +724,14 @@ def test_subclass_can_preserve_energy_precision(self):
         assert data.energy.dtype == torch.float64
         assert data.energy.item() == e64.item()
 
-    def test_non_exempt_fp_field_still_cast(self):
+    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 = _KeepEnergyData(
+            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),
diff --git a/test/data/test_batch.py b/test/data/test_batch.py
index 050d000f..b0657940 100644
--- a/test/data/test_batch.py
+++ b/test/data/test_batch.py
@@ -16,8 +16,6 @@
 
 from __future__ import annotations
 
-from typing import ClassVar
-
 import pytest
 import torch
 
@@ -26,12 +24,6 @@
 from nvalchemi.data.level_storage import MultiLevelStorage, UniformLevelStorage
 
 
-class _KeepEnergyData(AtomicData):
-    """AtomicData subclass that opts ``energy`` out of the fp-dtype downcast."""
-
-    _precision_preserving_keys: ClassVar[frozenset[str]] = frozenset({"energy"})
-
-
 def _minimal_atomic_data(
     num_nodes: int = 4,
     num_edges: int = 0,
@@ -1336,13 +1328,16 @@ class TestPrecisionPreservingBatch:
 
     ``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.
+    so the ``precision_preserving_keys`` exemption is honoured end-to-end.
     """
 
-    def test_fp64_energy_survives_from_data_list(self):
+    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 = [
-            _KeepEnergyData(
+            AtomicData(
                 positions=torch.randn(2, 3, dtype=torch.float32),
                 atomic_numbers=torch.ones(2, dtype=torch.long),
                 energy=e64.clone(),
@@ -1354,11 +1349,14 @@ def test_fp64_energy_survives_from_data_list(self):
         # 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):
+    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(
             [
-                _KeepEnergyData(
+                AtomicData(
                     positions=torch.randn(2, 3, dtype=torch.float32),
                     atomic_numbers=torch.ones(2, dtype=torch.long),
                     energy=e64.clone(),