Formation energies in ML potentials

[zulissimeta]

Summary of Changes

This PR adds support for formation energy calculations in MLP recipes via the FAIRChem FormationEnergyCalculator wrapper. Formation energies are now accessible through two new parameters in static_job() and relax_job():

• use_formation_energy (bool): Enables formation energy calculation (currently FAIRChem UMA with task_name='omat' only)
• formation_energy_kwargs (dict): Optional custom kwargs to pass to the FormationEnergyCalculator

Key Changes

1. Core Implementation (src/quacc/recipes/mlp/_base.py):

   • Added use_formation_energy and formation_energy_kwargs parameters to pick_calculator()
   • Implemented validation chain ensuring FAIRChem UMA omat configuration is used
   • Calculator wrapping handled transparently at the factory level, benefiting all recipes

2. API Exposure (src/quacc/recipes/mlp/core.py):

   • Added explicit parameters to static_job() and relax_job() for clarity and IDE autocomplete
   • Updated comprehensive docstrings

3. Testing (tests/core/recipes/mlp_recipes/test_core_recipes.py):

   • Added 6 new test functions covering:
     • Static calculations (elemental and binary compounds)
     • Relaxations (atomic and cell relaxation)
     • Multiple compound types
     • Error handling for unsupported configurations
   • All tests properly skip if fairchem not installed

4. Documentation (src/quacc/recipes/mlp/__init__.py):

   • Updated module docstring with FAIRChem-specific examples
   • Clarified limitations and supported configuration

Usage Example

from ase.build import bulk
from quacc.recipes.mlp.core import static_job

# Calculate formation energy for MgO
atoms = bulk("MgO", crystalstructure="rocksalt", a=4.2)
result = static_job(
    atoms,
    method="fairchem",
    name_or_path="uma-s-1",
    task_name="omat",
    use_formation_energy=True,
)
# result["results"]["energy"] contains formation energy (total energy, not per atom)

Requirements

• My PR is focused on a single feature addition or bugfix.
• My PR has relevant, comprehensive unit tests.
• My PR is on a custom branch (i.e. is not named main).

Note: If you are an external contributor, you will see a comment from @buildbot-princeton. This is solely for the maintainers.

[buildbot-princeton]

Can one of the admins verify this patch?

[codecov]

Codecov Report

❌ Patch coverage is 23.80952% with 32 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.88%. Comparing base (173dc82) to head (64c104e).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/quacc/recipes/mlp/_base.py	17.94%	32 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3122      +/-   ##
==========================================
- Coverage   97.69%   96.88%   -0.82%     
==========================================
  Files          97       97              
  Lines        4172     4210      +38     
==========================================
+ Hits         4076     4079       +3     
- Misses         96      131      +35     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Andrew-S-Rosen]

Hi @zulissimeta --- this looks like a nice addition to me. I have left a few comments.

I also think I ultimately need to make sure these fairchem tests run on our Princeton HPC clusters on PR because now that there is so much code, relying for it to run on main is going to be a problem long term. It is also difficult to review because the test coverage isn't reflective on reality. On our side, we'll work on doing that before the merge.

[Andrew-S-Rosen]

Is this something that is shipped with fairchem that we can access? I assume not but just want to check.

[Andrew-S-Rosen]

Many HPC clusters do not have access to the external network. Is there a mechanism we can encourage the users to call to run this in advance (e.g. on a login node)?

[Andrew-S-Rosen]

This stores it in the current working directory. It seems like this should instead be prefaced by Path(__file__).parent so it's stored in a fixed location.

Also, are we going to be redownloading the YAML every single time a job runs? That seems inefficient if so. It seems like we should be storing this once in and checking if it's there.

[Andrew-S-Rosen]

In here and the other function, do we need to keep these imports inside the function or can they be global?

[Andrew-S-Rosen]

Any idea why these test values for orb are being updated throughout the test suite?