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

[EricZQu]

ALCHEMI Toolkit Pull Request

Description

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

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Related Issues

Changes Made

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.

Testing

• Unit tests pass locally (make pytest)
• Linting passes (make lint)
• New tests added for new functionality meets coverage expectations?

Checklist

• I have read and understand the Contributing Guidelines
• I have updated the CHANGELOG.md
• I have performed a self-review of my code
• I have added docstrings to new functions/classes
• I have updated the documentation (if applicable)

Additional Notes

Tip

This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR adds an opt-in precision_preserving_keys ClassVar to AtomicData so callers can exempt specific fields (e.g. extensive total energy) from the check_fp_dtype_consistency cast to the positions dtype, and simultaneously reduces log noise by deduplicating the cast warning to once per (field, src_dtype, tgt_dtype) triple across the process lifetime.

• precision_preserving_keys exemption — a new empty frozenset ClassVar; setting it on a subclass or globally skips the fp-dtype cast for listed fields, preserving full precision through construction, add_system_property, batching, and device moves.
• Warning dedup — the new _warn_first helper and module-level _FP_CAST_WARNED set ensure each distinct cast combination warns only once, and the _reset_fp_cast_warned autouse fixture in conftest.py clears that set between tests to keep assertions order-independent.
• Tests — five new unit tests in TestDtypeCastWarning and a dedicated TestPrecisionPreservingBatch class cover the exemption, non-exempt casting, once-per-triple dedup, and the full batching/device-move path.

Important Files Changed

Filename	Overview
nvalchemi/data/atomic_data.py	Adds precision_preserving_keys ClassVar and _warn_first/_FP_CAST_WARNED for once-per-process cast warnings; logic is correct and well-structured
test/conftest.py	Adds autouse fixture to clear _FP_CAST_WARNED before each test; setup-only (no yield) but functionally sufficient for test isolation
test/data/test_atomic_data.py	Thorough new tests covering exemption, non-exempt casting, once-per-field-dtype warning dedup, and cross-field dedup
test/data/test_batch.py	New TestPrecisionPreservingBatch class verifies fp64 energy survives from_data_list, device moves, and the no-exemption baseline
CHANGELOG.md	Accurately documents the new opt-in precision-preserving feature and the warning dedup behavior

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into allow-fp64-ener..." | Re-trigger Greptile}

[EricZQu]

Added tests for batch, and a first time warning for casting

[laserkelvin]

I'll conditionally approve - I think I agree with the behavior change, but I'm not 100% set on how it's being done here.

I'm not 100% sure if this is best as part of the private API: the way you are using this in the test is by subclassing AtomicData which I think is not unreasonable, but I wonder if we should just let users be able to change the precision preservation without jumping through this hoop.

I'm considering these two scenarios:

• You modify the set once per Python session, so that the change in casting behavior is global - you can do this simply by making _precision_preserving_keys public
• You are working in a notebook environment, and you want to change the casting behavior temporarily for some subset of the data - for this you would need to have it as instance level

Can you consider the latter case? We can mark it out of scope for now, but I just want you to try and think of a good solution

[laserkelvin]

Ooof I'm not a big fan of this since it clutters up the schema...

I think I would be okay to have _FP_CAST_WARNED as like a module-level variable that tracks instead of it belonging to the schema (even though it's nominally a ClassVar).

[EricZQu]

Fixed!

[laserkelvin]

/ok to test ac8b6d2

[EricZQu]

Thanks a lot for the comments @laserkelvin! The API is public now and the warned set is now global. I thought about the per-instance "notebook" use case, and I think maybe we could use a context manager (with AtomicData.preserving({"energy"}): ...) for temporarily preserving a field without setting the global flag. It's doable but instance state isn't in model_dump, so a .to() device move drops the exemption. If we want the notebook ergonomics then we need to rewrite .to() for AtomicData. I could do it but idk if it's worth it.