Fix yaml serialization in phonon wf + add EPFL UPET models (#1400)

* safe load yaml in phonon calcs

* add upet models

* update mlff docs

* change default upet model to smallest omat model, test JSON store explicitly

Author: esoteric-ephemera

docs/user/codes/forcefields.md

@@ -36,6 +36,7 @@ Support is provided for the following models, which can be selected using `atoma
 | Neuroevolution Potential (NEP) | `NEP` | [10.1103/PhysRevB.104.104309](https://doi.org/10.1103/PhysRevB.104.104309) | Relies on `calorine` package |
 | Neural Equivariant Interatomic Potentials (Nequip) | `Nequip` | [10.1038/s41467-022-29939-5](https://doi.org/10.1038/s41467-022-29939-5) | Relies on the `nequip` package |
 | SevenNet | `SevenNet` | [10.1021/acs.jctc.4c00190](https://doi.org/10.1021/acs.jctc.4c00190) | Relies on the `sevenn` package |
+| Universal Point Edge Transformer (UPET) | `UPET` | [10.1038/s41467-025-65662-7](https://doi.org/10.1038/s41467-025-65662-7) | Relies on the `upet` package. Defaults to the smallest "pet-omad-xs" model. |
 
 ## Using custom forcefields by dictionary
 

pyproject.toml

@@ -64,6 +64,7 @@ forcefields = [
     "mattersim>=1.0.1",
     "sevenn>=0.9.3",
     "deepmd-kit>=2.1.4",
+    "upet>=0.2.0",
 ]
 torchsim = [
     "torch-sim-atomistic==0.5.0; python_version >= '3.12'"
@@ -102,6 +103,7 @@ strict-forcefields-generic = [
     "tensorflow-cpu==2.20.0",
     "mattersim==1.2.1",
     "wandb>=0.24.0", # required for mattersim
+    "upet>=0.2.0",
 ]
 strict-forcefields-torch-limited = [
     "matgl==2.0.6",

src/atomate2/common/schemas/phonons.py

@@ -16,7 +16,7 @@
 from pydantic import BaseModel, Field
 from pymatgen.core import Structure
 from pymatgen.io.phonopy import (
-    get_ph_bs_symm_line,
+    get_ph_bs_symm_line_from_dict,
     get_ph_dos,
     get_phonopy_structure,
     get_pmg_structure,
@@ -27,6 +27,7 @@
 from pymatgen.phonon.plotter import PhononBSPlotter, PhononDosPlotter
 from pymatgen.symmetry.bandstructure import HighSymmKpath
 from pymatgen.symmetry.kpath import KPathSeek
+from ruamel.yaml import YAML
 from typing_extensions import Self
 
 from atomate2.aims.utils.units import omegaToTHz
@@ -393,8 +394,10 @@ def from_forces_born(
             is_band_connection=kwargs.get("band_structure_eigenvectors", False),
         )
         phonon.write_yaml_band_structure(filename=filename_band_yaml)
-        bs_symm_line = get_ph_bs_symm_line(
-            filename_band_yaml, labels_dict=kpath_dict, has_nac=born is not None
+        bs_symm_line = get_ph_bs_symm_line_from_dict(
+            YAML(typ="safe").load(Path(filename_band_yaml).read_text()),
+            has_nac=born is not None,
+            labels_dict=kpath_dict,
         )
         new_plotter = PhononBSPlotter(bs=bs_symm_line)
         new_plotter.save_plot(

src/atomate2/forcefields/utils.py

@@ -45,6 +45,7 @@ class MLFF(Enum):  # TODO inherit from StrEnum when 3.11+
     Allegro = "Allegro"
     FAIRChem = "FAIRChem"
     MatterSim = "MatterSim"
+    UPET = "UPET"
 
     @classmethod
     def _missing_(cls, value: Any) -> Any:
@@ -80,6 +81,10 @@ def _missing_(cls, value: Any) -> Any:
         "predict_unit": {"model_name": "uma-s-1p1"},
         "task_name": "omat",
     },
+    MLFF.UPET: {
+        "model": "pet-omad-xs",
+        "version": "1.0.0",
+    },
 }
 
 
@@ -377,6 +382,11 @@ def ase_calculator(
 
                 calculator = MatterSimCalculator(**kwargs)
 
+            case MLFF.UPET:
+                from upet.calculator import UPETCalculator
+
+                calculator = UPETCalculator(**kwargs)
+
     elif isinstance(calculator_meta, dict):
         calc_cls = _load_calc_cls(calculator_meta)
         calculator = calc_cls(**kwargs)
@@ -444,6 +454,8 @@ def _get_pkg_name(calculator_meta: MLFF | dict[str, Any]) -> str | None:
                 ff_pkg = "calorine"
             case MLFF.SevenNet:
                 ff_pkg = "sevenn"
+            case MLFF.UPET:
+                ff_pkg = "upet"
             case _:
                 ff_pkg = None
         return ff_pkg

tests/forcefields/conftest.py

@@ -20,8 +20,12 @@
     try:
         _ = ase_calculator(mlff)
         _INSTALLED_MLFF[mlff.name] = True
-    except (ImportError, ValueError):
-        _INSTALLED_MLFF[mlff.name] = False
+    except (ImportError, ValueError) as exc:
+        # Note that UPET calculator raises a ValueError if no kwargs are passed
+        # Catch that here
+        _INSTALLED_MLFF[mlff.name] = (
+            "'model' parameter is required when not using checkpoint_path" in str(exc)
+        )
     except Exception:  # noqa: BLE001
         # Some calculators, like GAP, require extra potential files
         # Generally, thesea re

tests/forcefields/flows/test_phonon.py

@@ -1,9 +1,10 @@
 import os
 from pathlib import Path
+from tempfile import TemporaryDirectory
 
 import pytest
 from ase.calculators.calculator import Calculator
-from jobflow import Flow, run_locally
+from jobflow import Flow, JobStore, run_locally
 from numpy.testing import assert_allclose
 from pymatgen.core.structure import Structure
 from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine
@@ -113,6 +114,26 @@ def test_phonon_wf_force_field(
 ):
     # TODO brittle due to inability to adjust dtypes in CHGNetRelaxMaker
 
+    # See issue: https://github.com/materialsproject/atomate2/issues/1395
+    # Testing JSON store explicitly to avoid regression for this flow
+    json_dir = TemporaryDirectory()
+    json_store = JobStore.from_dict_spec(
+        {
+            "docs_store": {
+                "type": "JSONStore",
+                "paths": str(Path(json_dir.name) / "output.json"),
+                "read_only": False,
+            },
+            "additional_stores": {
+                "data": {
+                    "type": "JSONStore",
+                    "paths": str(Path(json_dir.name) / "blob_output.json"),
+                    "read_only": False,
+                }
+            },
+        }
+    )
+
     phonon_kwargs = dict(
         use_symmetrized_structure="conventional",
         create_thermal_displacements=False,
@@ -146,7 +167,12 @@ def test_phonon_wf_force_field(
     flow = phonon_maker.make(si_structure)
 
     # run the flow or job and ensure that it finished running successfully
-    responses = run_locally(flow, create_folders=True, ensure_success=True)
+    responses = run_locally(
+        flow, create_folders=True, ensure_success=True, store=json_store
+    )
+
+    # close temp dir for JSON store
+    json_dir.cleanup()
 
     # validate the outputs
     ph_bs_dos_doc = responses[flow[-1].uuid][1].output

tests/forcefields/test_jobs.py

@@ -785,6 +785,33 @@ def test_mattersim_relax_maker(si_structure: Structure, test_dir: Path):
     assert output.forcefield_version == get_imported_version("mattersim")
 
 
+@pytest.mark.skipif(not mlff_is_installed("UPET"), reason="upet is not installed.")
+def test_upet_relax_maker(si_structure: Structure, test_dir: Path):
+
+    # translate one atom to ensure a small number of relaxation steps are taken
+    si_structure.translate_sites(0, [0, 0, 0.1])
+    # generate job
+    job = ForceFieldRelaxMaker(
+        force_field_name="UPET",
+        steps=25,
+    ).make(si_structure)
+    responses = run_locally(job, ensure_success=True)
+    output = responses[job.uuid][1].output
+    assert isinstance(output, ForceFieldTaskDocument)
+    assert output.output.energy == approx(-11.241942405700684, rel=1e-4)
+    assert np.allclose(
+        output.output.ionic_steps[-1].forces,
+        [
+            [-0.16317394375801086, -0.11385675519704819, 0.09085404127836227],
+            [0.16317394375801086, 0.113856740295887, -0.09085404127836227],
+        ],
+        rtol=1e-2,
+    )
+    assert len(output.output.ionic_steps) > 1
+    assert output.output.n_steps == len(output.output.ionic_steps)
+    assert output.forcefield_version == get_imported_version("upet")
+
+
 @pytest.mark.skipif(not mlff_is_installed("MACE"), reason="mace_torch is not installed")
 def test_ext_load_static_maker(si_structure: Structure):
     calculator_meta = {

tests/forcefields/test_md.py

@@ -83,6 +83,7 @@ def test_ml_ff_md_maker(
         MLFF.MATPES_R2SCAN: -8.561729431152344,
         MLFF.FAIRChem: -5.4,
         MLFF.MatterSim: -5.4,
+        MLFF.UPET: -5.6,
     }
 
     # ASE can slightly change tolerances on structure positions