Formation energy calculator

docs/_toc.yml

@@ -72,7 +72,7 @@ parts:
   - file: inorganic_materials/models
   - file: inorganic_materials/examples_tutorials/summary
     sections:
-    - file: inorganic_materials/examples_tutorials/bulk_stability
+    - file: inorganic_materials/examples_tutorials/formation_energy
     - file: inorganic_materials/examples_tutorials/phonons
     - file: inorganic_materials/examples_tutorials/elastic
   - file: inorganic_materials/FAQ

docs/inorganic_materials/examples_tutorials/formation_energy.md

@@ -0,0 +1,94 @@
+---
+jupytext:
+  text_representation:
+    extension: .md
+    format_name: myst
+    format_version: 0.13
+    jupytext_version: 1.17.1
+kernelspec:
+  display_name: Python 3 (ipykernel)
+  language: python
+  name: python3
+---
+
+Formation energies
+------------------
+
+We're going to start simple here - let's run a local relaxation (optimize the unit cell and positions) using a pre-trained UMA model to compute formation energies for inorganic materials.
+
+Note predicting formation energy using models that models trained solely on OMat24 must use OMat24 compatible references and corrections for mixing PBE and PBE+U calculations. We use MP2020-style corrections fitted to OMat24 DFT calculations. For more information see the [documentation](https://docs.materialsproject.org/methodology/materials-methodology/thermodynamic-stability/thermodynamic-stability/anion-and-gga-gga+u-mixing) at the Materials Project. The necessary references can be found using the `fairchem.data.omat` package!
+
+````{admonition} Need to install fairchem-core or get UMA access or getting permissions/401 errors?
+:class: dropdown
+
+
+1. Install the necessary packages using pip, uv etc
+```{code-cell} ipython3
+:tags: [skip-execution]
+
+! pip install fairchem-core fairchem-data-omat
+```
+
+2. Get access to any necessary huggingface gated models
+    * Get and login to your Huggingface account
+    * Request access to https://huggingface.co/facebook/UMA
+    * Create a Huggingface token at https://huggingface.co/settings/tokens/ with the permission "Permissions: Read access to contents of all public gated repos you can access"
+    * Add the token as an environment variable using `huggingface-cli login` or by setting the HF_TOKEN environment variable.
+
+```{code-cell} ipython3
+:tags: [skip-execution]
+
+# Login using the huggingface-cli utility
+! huggingface-cli login
+
+# alternatively,
+import os
+os.environ['HF_TOKEN'] = 'MY_TOKEN'
+```
+````
+
+```{code-cell} ipython3
+from __future__ import annotations
+
+import pprint
+
+from ase.build import bulk
+from ase.optimize import LBFGS
+from quacc.recipes.mlp.core import relax_job
+
+from fairchem.core.calculate.ase_calculator import FAIRChemCalculator, set_predict_formation_energy
+
+# Make an Atoms object of a bulk MgO structure
+atoms = bulk("MgO", "rocksalt", a=4.213)
+
+# Run a structure relaxation
+result = relax_job(
+    atoms,
+    method="fairchem",
+    name_or_path="uma-s-1p1",
+    task_name="omat",
+    relax_cell=True,
+    opt_params={"fmax": 1e-3, "optimizer": FIRE},
+)
+
+# Get the realxed atoms!
+atoms = result["atoms"]
+
+# Create an calculator using uma-s-1p1
+calculator = FAIRChemCalculator.from_model_checkpoint("uma-s-1p1", task_name="omat")
+atoms.calc = calculator
+
+# Adapt the calculation to automatically return MP-style corrected formation energies
+# For the omat task, this defaults to apply MP2020 style corrections with OMat24 compatibility
+with set_predict_formation_energy(atoms.calc, apply_corrections=True):
+    form_energy = atoms.get_potential_energy()
+```
+
+```{code-cell} ipython3
+pprint.pprint(f"Total energy: {result["results"]["energy"] eV\n Formation energy {form_energy} eV})
+```
+
+Compare the results to the value of [-3.038 eV/atom reported](https://next-gen.materialsproject.org/materials/mp-1265?chemsys=Mg-O#thermodynamic_stability) in the the Materials Project!
+*Note that we expect differences due to the different DFT settings used to calculate the OMat24 training data.*
+
+Congratulations; you ran your first relaxation and predicted the formation energy of MgO using UMA and `quacc`!

src/fairchem/core/calculate/ase_calculator.py

@@ -9,6 +9,8 @@
 
 import logging
 import os
+from collections import Counter
+from contextlib import contextmanager
 from functools import partial
 from typing import TYPE_CHECKING, Literal
 
@@ -316,6 +318,89 @@ def _validate_charge_and_spin(self, atoms: Atoms) -> None:
             )
 
 
+@contextmanager
+def set_predict_formation_energy(
+    calculator: FAIRChemCalculator,
+    element_references: dict | None = None,
+    apply_corrections: bool | None = None,
+    clear_calculator_cache: bool = True,
+) -> FAIRChemCalculator:
+    """
+    Adapt a calculator to predict formation energy.
+
+    Args:
+        calculator (FAIRChemCalculator): The calculator instance to modify.
+        element_references (dict, optional): Optional dictionary of formation reference energies for each element. You likely do not want
+            to provide these and instead use the defaults for each UMA task.
+        apply_corrections (bool, optional): Whether to apply MP style corrections to the formation energies.
+            This is only relevant for the OMat task. Default is True if task is OMat.
+        clear_calculator_cache (bool): Whether to clear the calculator cache before modifying the calculate method.
+
+    Returns:
+        FAIRChemCalculator: The same calculator instance but will return formation energies as the potential energy.
+    """
+    if element_references is None:
+        element_references = calculator.predictor.form_elem_refs[calculator.task_name]
+
+    if apply_corrections is True and calculator.task_name != UMATask.OMAT.value:
+        raise ValueError("MP style corrections can only be applied for the OMat task.")
+    if apply_corrections is None and calculator.task_name == UMATask.OMAT.value:
+        apply_corrections = True
+
+    original_calculate = calculator.calculate
+
+    if clear_calculator_cache is True:
+        calculator.atoms = None
+
+    def formation_energy_calculate(
+        atoms: Atoms, properties: list[str], system_changes: list[str]
+    ) -> None:
+        original_calculate(atoms, properties, system_changes)
+
+        if "energy" in calculator.results:
+            total_energy = calculator.results["energy"]
+
+            if apply_corrections:
+                try:
+                    from fairchem.data.omat.entries.compatibility import (
+                        apply_mp_style_corrections,
+                    )
+                except ImportError as err:
+                    raise ImportError(
+                        "fairchem.data.omat is required to apply MP style corrections. Please install it."
+                    ) from err
+                total_energy = apply_mp_style_corrections(total_energy, atoms)
+
+            element_symbols = atoms.get_chemical_symbols()
+            element_counts = Counter(element_symbols)
+
+            missing_elements = set(element_symbols) - set(element_references.keys())
+            if missing_elements:
+                raise ValueError(
+                    f"Missing reference energies for elements: {missing_elements}"
+                )
+
+            total_ref_energy = sum(
+                element_references[element] * count
+                for element, count in element_counts.items()
+            )
+
+            formation_energy = total_energy - total_ref_energy
+
+            calculator.results["energy"] = formation_energy
+
+            if "free_energy" in calculator.results:
+                calculator.results["free_energy"] = formation_energy
+
+    try:
+        calculator.calculate = formation_energy_calculate
+        yield
+    finally:
+        calculator.calculate = original_calculate
+        if clear_calculator_cache is True:
+            calculator.atoms = None
+
+
 class MixedPBCError(ValueError):
     """Specific exception example."""
 

src/fairchem/core/calculate/pretrained_mlip.py

@@ -30,6 +30,7 @@ class HuggingFaceCheckpoint:
     subfolder: str | None = None  # specify a hf repo subfolder
     revision: str | None = None  # specify a version tag, branch, commit hash
     atom_refs: dict | None = None  # specify an isolated atomic reference
+    form_elem_refs: dict | None = None  # specify a form elemental reference
 
 
 @dataclass
@@ -94,31 +95,51 @@ def get_predict_unit(
         KeyError: If the specified model_name is not found in available models.
     """
     checkpoint_path = pretrained_checkpoint_path_from_name(model_name)
-    atom_refs = get_isolated_atomic_energies(model_name, cache_dir)
+    atom_refs = get_reference_energies(model_name, "atom_refs", cache_dir)
+
+    if _MODEL_CKPTS.checkpoints[model_name].form_elem_refs is not None:
+        form_elem_refs = get_reference_energies(
+            model_name, "form_elem_refs", cache_dir
+        )["refs"]
+    else:
+        form_elem_refs = None
+
     return load_predict_unit(
-        checkpoint_path, inference_settings, overrides, device, atom_refs, workers
+        checkpoint_path,
+        inference_settings,
+        overrides,
+        device,
+        atom_refs,
+        form_elem_refs,
+        workers,
     )
 
 
-def get_isolated_atomic_energies(model_name: str, cache_dir: str = CACHE_DIR) -> dict:
+def get_reference_energies(
+    model_name: str,
+    reference_type: Literal["atom_refs", "form_elem_refs"] = "atom_refs",
+    cache_dir: str = CACHE_DIR,
+) -> dict:
     """
     Retrieves the isolated atomic energies for use with single atom systems into the CACHE_DIR
 
     Args:
         model_name: Name of the model to load from available pretrained models.
+        reference_type: Type of references file to download: atom_refs or bulk_refs.
         cache_dir: Path to folder where files will be stored. Default is "~/.cache/fairchem"
     Returns:
-        Atomic element reference data
+        Atomic or bulk phase element reference data
 
     Raises:
         KeyError: If the specified model_name is not found in available models.
     """
     model_checkpoint = _MODEL_CKPTS.checkpoints[model_name]
-    atomic_refs_path = hf_hub_download(
-        filename=model_checkpoint.atom_refs["filename"],
+    file_data = getattr(model_checkpoint, reference_type)
+    refs_path = hf_hub_download(
+        filename=file_data["filename"],
         repo_id=model_checkpoint.repo_id,
-        subfolder=model_checkpoint.atom_refs["subfolder"],
+        subfolder=file_data["subfolder"],
         revision=model_checkpoint.revision,
         cache_dir=cache_dir,
     )
-    return OmegaConf.load(atomic_refs_path)
+    return OmegaConf.load(refs_path)

src/fairchem/core/calculate/pretrained_models.json

@@ -6,6 +6,10 @@
         "atom_refs": {
             "subfolder": "references",
             "filename": "iso_atom_elem_refs.yaml"
+        },
+        "form_elem_refs": {
+            "subfolder": "references",
+            "filename": "form_elem_refs.yaml"
         }
     },
     "uma-s-1p1": {
@@ -15,6 +19,10 @@
         "atom_refs": {
             "subfolder": "references",
             "filename": "iso_atom_elem_refs.yaml"
+        },
+        "form_elem_refs": {
+            "subfolder": "references",
+            "filename": "form_elem_refs.yaml"
         }
     },
     "uma-m-1p1": {
@@ -24,6 +32,10 @@
         "atom_refs": {
             "subfolder": "references",
             "filename": "iso_atom_elem_refs.yaml"
+        },
+        "form_elem_refs": {
+            "subfolder": "references",
+            "filename": "form_elem_refs.yaml"
         }
     },
     "esen-md-direct-all-omol": {