Skip to content

Avoid recomputing the model on every property read in FormationEnergyCalculator - #2212

Open
u7k4rs6 wants to merge 2 commits into
facebookresearch:mainfrom
u7k4rs6:cache-formation-energy-results
Open

u7k4rs6 wants to merge 2 commits into
facebookresearch:mainfrom
u7k4rs6:cache-formation-energy-results

Conversation

@u7k4rs6

@u7k4rs6 u7k4rs6 commented Sep 30, 2026

Copy link
Copy Markdown

FormationEnergyCalculator.calculate() calls the wrapped calculator directly but never calls Calculator.calculate(self, ...), so the wrapper's self.atoms is never set. ASE's check_state then reports changes on every access, and every get_potential_energy / get_forces / get_stress call reruns the full model even when nothing changed.

In practice this hurts relaxations and MD a lot, since optimizers read energy and forces separately. With a call-counting EMT wrapped in FormationEnergyCalculator, a 10-step BFGS run did 44 base calculations on main vs 11 with this change.

The fix calls Calculator.calculate(self, ...) first so ASE can cache, and copies the check_state override from FAIRChemCalculator so changing atoms.info (charge/spin) still triggers a recompute.

Added test_formation_energy_calculator_caches_results. It uses a small counting calculator (no checkpoint needed) and checks that repeated reads hit the base calculator once, and that moving atoms or changing atoms.info recomputes. It fails on main and passes here. ruff (pinned 0.5.1) is clean.

@meta-cla meta-cla Bot added the cla signed label Sep 30, 2026
@zulissimeta zulissimeta added enhancement New feature or request patch Patch version release labels Sep 30, 2026
@zulissimeta
zulissimeta self-requested a review September 30, 2026 20:47

@zulissimeta zulissimeta left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great catch! Thanks for the PR.

This branch has not been deployed

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

Labels

cla signed enhancement New feature or request patch Patch version release

2 participants