Skip to content

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

Open
EricZQu wants to merge 2 commits into
NVIDIA:mainfrom
EricZQu:allow-fp64-energy
Open

Add opt-in hook to exempt fields from the fp-dtype downcast#127
EricZQu wants to merge 2 commits into
NVIDIA:mainfrom
EricZQu:allow-fp64-energy

Conversation

@EricZQu

@EricZQu EricZQu commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

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.

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>
@copy-pr-bot

copy-pr-bot Bot commented Jul 3, 2026

Copy link
Copy Markdown

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-apps Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds an opt-in _precision_preserving_keys ClassVar to AtomicData that exempts specified floating-point fields from the positions-dtype downcast in check_fp_dtype_consistency. The default is an empty frozenset, so existing behavior is unchanged.

  • The exemption is correctly applied before the getattr call and works for both declared fields and Pydantic extra fields, covering dynamically added properties via add_system_property / add_node_property since validate_assignment=True re-runs the validator on each setattr.
  • Three focused tests cover the default downcast behavior, the subclass exemption (including bit-exact precision verification and re-validation persistence), and the invariant that non-exempt FP fields are still cast and warned about.

Important Files Changed

Filename Overview
nvalchemi/data/atomic_data.py Adds _precision_preserving_keys ClassVar and a two-line skip in check_fp_dtype_consistency; logic is correct and the early-continue is placed before getattr to avoid unnecessary attribute access on exempt keys.
test/data/test_atomic_data.py Adds a _KeepEnergyData test subclass and three tests covering default downcast, precision preservation with bit-exact assertion, re-validation persistence after add_system_property, and non-exempt fields still being cast and warned.
CHANGELOG.md Changelog entry clearly documents the new ClassVar, its default, and the motivating precision problem with extensive energies.

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant