Bump chgnet from 0.3.8 to 0.4.0 (#990)

* Bump chgnet from 0.3.8 to 0.4.0

Bumps [chgnet](https://github.com/CederGroupHub/chgnet) from 0.3.8 to 0.4.0.
- [Release notes](https://github.com/CederGroupHub/chgnet/releases)
- [Commits](CederGroupHub/chgnet@v0.3.8...v0.4.0)

---
updated-dependencies:
- dependency-name: chgnet
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <[email protected]>

* strict-forcefields bump torch==2.2.1 -> 2.4.1 for chgnet compat

---------

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Janosh Riebesell <[email protected]>

Author: dependabot[bot]

pyproject.toml

@@ -118,12 +118,12 @@ strict = [
 ]
 strict-forcefields = [
     "calorine==2.2.1",
-    "chgnet==0.3.8",
+    "chgnet==0.4.0",
     "mace-torch>=0.3.3",
     "matgl==1.1.3",
     "quippy-ase==0.9.14; python_version < '3.12'",
     "sevenn==0.9.3.post1",
-    "torch==2.2.1",
+    "torch==2.4.1",
     "torchdata==0.7.1",                            # TODO: remove when issue fixed
 ]
 

src/atomate2/forcefields/utils.py

@@ -42,7 +42,7 @@ def ase_calculator(calculator_meta: str | dict, **kwargs: Any) -> Calculator | N
     """
     calculator = None
 
-    if isinstance(calculator_meta, str) and calculator_meta in map(str, MLFF):
+    if isinstance(calculator_meta, str | MLFF) and calculator_meta in map(str, MLFF):
         calculator_name = MLFF(calculator_meta.split("MLFF.")[-1])
 
         if calculator_name == MLFF.CHGNet:
@@ -88,7 +88,7 @@ def ase_calculator(calculator_meta: str | dict, **kwargs: Any) -> Calculator | N
         calculator = calc_cls(**kwargs)
 
     if calculator is None:
-        raise ValueError("Could not create ASE calculator.")
+        raise ValueError(f"Could not create ASE calculator for {calculator_meta}.")
 
     return calculator
 

src/atomate2/qchem/drones.py

@@ -1,4 +1,4 @@
-"""Drones for parsing VASP calculations and realtd outputs."""
+"""Drones for parsing VASP calculations and related outputs."""
 
 from __future__ import annotations
 
@@ -69,15 +69,15 @@ def get_valid_paths(self, path: tuple[str, list[str], list[str]]) -> list[str]:
         parent, subdirs, _ = path
         task_names = ["mol.qout.*"]
         combined_paths = [parent + os.sep + sdir for sdir in subdirs]
-        rpath = []
-        for cpath in combined_paths:
-            fnames = os.listdir(cpath)
-            if any(name.startswith("mol.qout.") for name in fnames):
-                rpath.append(parent)
+        valid_paths = []
+        for sub_dir in combined_paths:
+            file_names = os.listdir(sub_dir)
+            if any(name.startswith("mol.qout.") for name in file_names):
+                valid_paths.append(parent)
 
             if (
                 not any(parent.endswith(os.sep + r) for r in task_names)
                 and len(list(Path(parent).glob("mol.qout*"))) > 0
             ):
-                rpath.append(parent)
-        return rpath
+                valid_paths.append(parent)
+        return valid_paths

tests/aims/test_flows/test_anharmonic_quantification.py

@@ -60,10 +60,7 @@ def test_anharmonic_quantification_oneshot(si, clean_dir, mock_aims, species_dir
         phonon_maker=phonon_maker,
     )
     maker.name = "anharmonicity"
-    flow = maker.make(
-        si,
-        supercell_matrix=np.array([-1, 1, 1, 1, -1, 1, 1, 1, -1]).reshape((3, 3)),
-    )
+    flow = maker.make(si, supercell_matrix=np.ones((3, 3)) - 2 * np.eye(3))
 
     # run the flow or job and ensure that it finished running successfully
     responses = run_locally(flow, create_folders=True, ensure_success=True)
@@ -120,7 +117,7 @@ def test_anharmonic_quantification_full(si, clean_dir, mock_aims, species_dir):
     maker.name = "anharmonicity"
     flow = maker.make(
         si,
-        supercell_matrix=np.array([-1, 1, 1, 1, -1, 1, 1, 1, -1]).reshape((3, 3)),
+        supercell_matrix=np.ones((3, 3)) - 2 * np.eye(3),
         one_shot_approx=False,
         seed=1234,
     )
@@ -254,9 +251,7 @@ def test_site_resolved_anharmonic_quantification(
     )
     maker.name = "anharmonicity"
     flow = maker.make(
-        nacl,
-        supercell_matrix=np.array([-1, 1, 1, 1, -1, 1, 1, 1, -1]).reshape((3, 3)),
-        site_resolved=True,
+        nacl, supercell_matrix=np.ones((3, 3)) - 2 * np.eye(3), site_resolved=True
     )
 
     # run the flow or job and ensure that it finished running successfully
@@ -323,7 +318,7 @@ def test_element_resolved_anharmonic_quantification(
     maker.name = "anharmonicity"
     flow = maker.make(
         nacl,
-        supercell_matrix=np.array([-1, 1, 1, 1, -1, 1, 1, 1, -1]).reshape((3, 3)),
+        supercell_matrix=np.ones((3, 3)) - 2 * np.eye(3),
         element_resolved=True,
     )
 
@@ -372,10 +367,7 @@ def test_anharmonic_quantification_socket_oneshot(si, clean_dir, species_dir):
         phonon_maker=phonon_maker,
     )
     maker.name = "anharmonicity"
-    flow = maker.make(
-        si,
-        supercell_matrix=np.array([-1, 1, 1, 1, -1, 1, 1, 1, -1]).reshape((3, 3)),
-    )
+    flow = maker.make(si, supercell_matrix=np.ones((3, 3)) - 2 * np.eye(3))
 
     # run the flow or job and ensure that it finished running successfully
     responses = run_locally(flow, create_folders=True, ensure_success=True)

tests/aims/test_flows/test_phonon_workflow.py

@@ -48,10 +48,7 @@ def test_phonon_flow(si, clean_dir, mock_aims, species_dir):
         ),
     )
     maker.name = "phonons"
-    flow = maker.make(
-        si,
-        supercell_matrix=np.array([-1, 1, 1, 1, -1, 1, 1, 1, -1]).reshape((3, 3)),
-    )
+    flow = maker.make(si, supercell_matrix=np.ones((3, 3)) - 2 * np.eye(3))
 
     # run the flow or job and ensure that it finished running successfully
     responses = run_locally(flow, create_folders=True, ensure_success=True)
@@ -142,10 +139,7 @@ def test_phonon_socket_flow(si, clean_dir, mock_aims, species_dir):
         ),
     )
     maker.name = "phonons"
-    flow = maker.make(
-        si,
-        supercell_matrix=np.array([-1, 1, 1, 1, -1, 1, 1, 1, -1]).reshape((3, 3)),
-    )
+    flow = maker.make(si, supercell_matrix=np.ones((3, 3)) - 2 * np.eye(3))
 
     # run the flow or job and ensure that it finished running successfully
     responses = run_locally(flow, create_folders=True, ensure_success=True)
@@ -212,10 +206,8 @@ def test_phonon_default_flow(si, clean_dir, mock_aims, species_dir):
 
     maker = PhononMaker()
     maker.name = "phonons"
-    flow = maker.make(
-        si,
-        supercell_matrix=np.array([-1, 1, 1, 1, -1, 1, 1, 1, -1]).reshape((3, 3)),
-    )
+    supercell_matrix = np.ones((3, 3)) - 2 * np.eye(3)
+    flow = maker.make(si, supercell_matrix=supercell_matrix)
 
     # run the flow or job and ensure that it finished running successfully
     responses = run_locally(flow, create_folders=True, ensure_success=True)
@@ -288,10 +280,7 @@ def test_phonon_default_socket_flow(si, clean_dir, mock_aims, species_dir):
 
     maker = PhononMaker(socket=True)
     maker.name = "phonons"
-    flow = maker.make(
-        si,
-        supercell_matrix=np.array([-1, 1, 1, 1, -1, 1, 1, 1, -1]).reshape((3, 3)),
-    )
+    flow = maker.make(si, supercell_matrix=np.ones((3, 3)) - 2 * np.eye(3))
 
     # run the flow or job and ensure that it finished running successfully
     responses = run_locally(flow, create_folders=True, ensure_success=True)

tests/ase/test_jobs.py

@@ -33,6 +33,9 @@ def test_lennard_jones_relax_maker(lj_fcc_ne_pars, fcc_ne_structure):
     assert output.structure.volume == pytest.approx(22.304245)
     assert output.output.energy == pytest.approx(-0.018494767)
     assert isinstance(output, AseStructureTaskDoc)
+    assert fcc_ne_structure.matches(
+        output.structure
+    ), f"{output.structure} != {fcc_ne_structure}"
 
 
 def test_lennard_jones_static_maker(lj_fcc_ne_pars, fcc_ne_structure):
@@ -46,8 +49,12 @@ def test_lennard_jones_static_maker(lj_fcc_ne_pars, fcc_ne_structure):
     assert output.structure.volume == pytest.approx(24.334)
     assert isinstance(output, AseStructureTaskDoc)
 
+    # Structure.__eq__ checks properties which contains 'energy', 'forces', 'stress'
+    # so need to reset properties to ensure equality
     output.structure.properties = fcc_ne_structure.properties
-    assert output.structure == fcc_ne_structure
+    assert (
+        output.structure == fcc_ne_structure
+    ), f"{output.structure} != {fcc_ne_structure}"
 
 
 @pytest.mark.skipif(condition=TBLite is None, reason="TBLite must be installed.")
@@ -94,9 +101,7 @@ def test_gfn_xtb_relax_maker(h2o_3uud_trimer):
 def test_gfn_xtb_static_maker(h2o_3uud_trimer):
     os.environ["OMP_NUM_THREADS"] = "1"
     job = GFNxTBStaticMaker(
-        calculator_kwargs={
-            "method": "GFN2-xTB",
-        },
+        calculator_kwargs={"method": "GFN2-xTB"},
     ).make(h2o_3uud_trimer)
 
     response = run_locally(job)
@@ -105,5 +110,7 @@ def test_gfn_xtb_static_maker(h2o_3uud_trimer):
     assert output.output.energy_per_atom == pytest.approx(-46.05920227158222)
     assert isinstance(output, AseMoleculeTaskDoc)
 
+    # Molecule.__eq__ checks properties which contains 'energy', 'forces', 'stress'
+    # so need to reset properties to ensure equality
     output.molecule.properties = h2o_3uud_trimer.properties
     assert output.molecule == h2o_3uud_trimer

tests/forcefields/flows/test_eos.py

@@ -6,13 +6,15 @@
 from atomate2.forcefields.flows.eos import (
     CHGNetEosMaker,
     ForceFieldEosMaker,
-    M3GNetEosMaker,
+    # M3GNetEosMaker,
     MACEEosMaker,
 )
 
 ff_maker_map = {
     MLFF.CHGNet.value: CHGNetEosMaker,
-    MLFF.M3GNet.value: M3GNetEosMaker,
+    # skip m3gnet due M3GNet requiring DGL which is PyTorch 2.4 incompatible
+    # raises "FileNotFoundError: Cannot find DGL C++ libgraphbolt_pytorch_2.4.1.so"
+    # MLFF.M3GNet.value: M3GNetEosMaker,
     MLFF.MACE.value: MACEEosMaker,
 }
 

tests/forcefields/test_jobs.py

@@ -132,6 +132,7 @@ def test_chgnet_relax_maker(si_structure: Structure, relax_cell: bool):
         CHGNetRelaxMaker()
 
 
+@pytest.mark.skip(reason="M3GNet requires DGL which is PyTorch 2.4 incompatible")
 def test_m3gnet_static_maker(si_structure):
     # generate job
     job = ForceFieldStaticMaker(
@@ -154,6 +155,7 @@ def test_m3gnet_static_maker(si_structure):
         M3GNetStaticMaker()
 
 
+@pytest.mark.skip(reason="M3GNet requires DGL which is PyTorch 2.4 incompatible")
 def test_m3gnet_relax_maker(si_structure):
     # translate one atom to ensure a small number of relaxation steps are taken
     si_structure.translate_sites(0, [0, 0, 0.1])

tests/forcefields/test_md.py

@@ -13,6 +13,7 @@
 from pymatgen.analysis.structure_matcher import StructureMatcher
 from pymatgen.core import Structure
 
+from atomate2.forcefields import MLFF
 from atomate2.forcefields.md import (
     CHGNetMDMaker,
     ForceFieldMDMaker,
@@ -24,12 +25,12 @@
 )
 
 name_to_maker = {
-    "CHGNet": CHGNetMDMaker,
-    "M3GNet": M3GNetMDMaker,
-    "MACE": MACEMDMaker,
-    "GAP": GAPMDMaker,
-    "NEP": NEPMDMaker,
-    "Nequip": NequipMDMaker,
+    MLFF.CHGNet: CHGNetMDMaker,
+    MLFF.M3GNet: M3GNetMDMaker,
+    MLFF.MACE: MACEMDMaker,
+    MLFF.GAP: GAPMDMaker,
+    MLFF.NEP: NEPMDMaker,
+    MLFF.Nequip: NequipMDMaker,
 }
 
 
@@ -47,40 +48,42 @@ def test_maker_initialization():
         )
 
 
-@pytest.mark.parametrize(
-    "ff_name",
-    ["CHGNet", "M3GNet", "MACE", "GAP", "NEP", "Nequip"],
-)
+@pytest.mark.parametrize("ff_name", MLFF)
 def test_ml_ff_md_maker(
     ff_name, si_structure, sr_ti_o3_structure, al2_au_structure, test_dir, clean_dir
 ):
-    if ff_name == "GAP" and sys.version_info >= (3, 12):
+    if ff_name == MLFF.Forcefield:
+        return  # nothing to test here, MLFF.Forcefield is just a generic placeholder
+    if ff_name == MLFF.GAP and sys.version_info >= (3, 12):
         pytest.skip(
             "GAP model not compatible with Python 3.12, waiting on https://github.com/libAtoms/QUIP/issues/645"
         )
+    if ff_name == MLFF.M3GNet:
+        pytest.skip("M3GNet requires DGL which is PyTorch 2.4 incompatible")
 
     n_steps = 5
 
     ref_energies_per_atom = {
-        "CHGNet": -5.280157089233398,
-        "M3GNet": -5.387282371520996,
-        "MACE": -5.311369895935059,
-        "GAP": -5.391255755606209,
-        "NEP": -3.966232215741286,
-        "Nequip": -8.84670181274414,
+        MLFF.CHGNet: -5.280157089233398,
+        MLFF.M3GNet: -5.387282371520996,
+        MLFF.MACE: -5.311369895935059,
+        MLFF.GAP: -5.391255755606209,
+        MLFF.NEP: -3.966232215741286,
+        MLFF.Nequip: -8.84670181274414,
+        MLFF.SevenNet: -5.394115447998047,
     }
 
     # ASE can slightly change tolerances on structure positions
     matcher = StructureMatcher()
 
     calculator_kwargs = {}
     unit_cell_structure = si_structure.copy()
-    if ff_name == "GAP":
+    if ff_name == MLFF.GAP:
         calculator_kwargs = {
             "args_str": "IP GAP",
             "param_filename": str(test_dir / "forcefields" / "gap" / "gap_file.xml"),
         }
-    elif ff_name == "NEP":
+    elif ff_name == MLFF.NEP:
         # NOTE: The test NEP model is specifically trained on 16 elemental metals
         # thus a new Al2Au structure is added.
         # The NEP model used for the tests is licensed under a
@@ -91,7 +94,7 @@ def test_ml_ff_md_maker(
             "model_filename": test_dir / "forcefields" / "nep" / "nep.txt"
         }
         unit_cell_structure = al2_au_structure.copy()
-    elif ff_name == "Nequip":
+    elif ff_name == MLFF.Nequip:
         calculator_kwargs = {
             "model_path": test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth"
         }
@@ -137,9 +140,9 @@ def test_ml_ff_md_maker(
         for key in ("energy", "forces", "stress", "velocities", "temperature")
         for step in task_doc.objects["trajectory"].frame_properties
     )
-
-    with pytest.warns(FutureWarning):
-        name_to_maker[ff_name]()
+    if ff_maker := name_to_maker.get(ff_name):
+        with pytest.warns(FutureWarning):
+            ff_maker()
 
 
 @pytest.mark.parametrize("traj_file", ["trajectory.json.gz", "atoms.traj"])

tests/forcefields/test_phonon.py

@@ -64,6 +64,11 @@ def test_phonon_maker_initialization_with_all_mlff(
         if mlff in {MLFF.GAP, MLFF.Forcefield}:
             continue  # TODO fix GAP, currently fails with RuntimeError, see
             # https://github.com/materialsproject/atomate2/pull/918#issuecomment-2253659694
+        # skip m3gnet due M3GNet requiring DGL which is PyTorch 2.4 incompatible
+        # raises "FileNotFoundError: Cannot find DGL C++ libgraphbolt_pytorch_2.4.1.so"
+        if mlff == MLFF.M3GNet:
+            continue
+
         calc_kwargs = {
             MLFF.Nequip: {"model_path": f"{chk_pt_dir}/nequip/nequip_ff_sr_ti_o3.pth"},
             MLFF.NEP: {"model_filename": f"{test_dir}/forcefields/nep/nep.txt"},