diff --git a/CHANGELOG.md b/CHANGELOG.md index 05bea77e..b3443bc6 100644 --- a/CHANGELOG.md +++ b/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[]` 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 diff --git a/nvalchemi/data/atomic_data.py b/nvalchemi/data/atomic_data.py index 1bf66d62..36504e2b 100644 --- a/nvalchemi/data/atomic_data.py +++ b/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") diff --git a/test/conftest.py b/test/conftest.py index d3647882..f95f0d09 100644 --- a/test/conftest.py +++ b/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: diff --git a/test/data/test_atomic_data.py b/test/data/test_atomic_data.py index 6c8dc10c..b085efd0 100644 --- a/test/data/test_atomic_data.py +++ b/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 diff --git a/test/data/test_batch.py b/test/data/test_batch.py index 34b62e19..b0657940 100644 --- a/test/data/test_batch.py +++ b/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