Stabilze forcefield dependencies (#1452)

* opt in forcefields

* minimal forcefield deps, use chgnet package for this + docs

* other test data

Author: esoteric-ephemera

.github/workflows/testing.yml

@@ -62,9 +62,9 @@ jobs:
           micromamba install -n a2 -c conda-forge packmol --yes
 
       - name: Install dependencies
-        run: | # TODO: remove setuptools pin once `mattersim` removes `pkg_resources` dependency
+        run: |
           micromamba activate a2
-          python -m pip install --upgrade pip 'setuptools<81'
+          python -m pip install --upgrade pip
           mkdir -p ~/.abinit/pseudos
           cp -r tests/test_data/abinit/pseudos/ONCVPSP-PBE-SR-PDv0.4 ~/.abinit/pseudos
           uv pip install .[strict,strict-forcefields-${{ matrix.dep-group }},abinit,approxneb,aims] --group tests
@@ -207,8 +207,8 @@ jobs:
         run: micromamba run -n a2 pip install uv
 
       - name: Install conda dependencies
-        run: |  # TODO: migrate openff tests to use non smirnoff99frosst forcefields - requires old setuptools / pkg_resources
-          micromamba install -n a2 -c conda-forge enumlib packmol bader openbabel openff-toolkit==0.16.2 openff-interchange==0.3.22 'setuptools<81' --yes
+        run: |  # TODO: migrate openff tests to use non smirnoff99frosst forcefields as recommended by devs
+          micromamba install -n a2 -c conda-forge enumlib packmol bader openbabel openff-toolkit==0.16.2 openff-interchange==0.3.22 --yes
 
       - name: Install dependencies
         run: |
@@ -376,8 +376,8 @@ jobs:
           cache-dependency-path: pyproject.toml
 
       - name: Install dependencies
-        run: | # TODO: remove setuptools pin once `mattersim` removes `pkg_resources` dependency
-          python -m pip install --upgrade pip 'setuptools<81'
+        run: |
+          python -m pip install --upgrade pip
           pip install .[strict,strict-forcefields-generic] --group docs
 
       - name: Build

docs/user/codes/forcefields.md

@@ -4,6 +4,13 @@
 
 `atomate2` includes an interface to a few common machine learning interatomic potentials (MLIPs), also known variously as machine learning forcefields (MLFFs), or foundation potentials (FPs) for universal variants.
 
+***As of `atomate2==0.1.1`, most forcefield packages are opt-in only. You must install those forcefields which you plan to use.***
+Running `pip install 'atomate2[forcefields]'` will install the `chgnet` package to permit you to try the forcefield classes.
+You can then select other forcefields you want to use.
+
+We have made this choice both to avoid the appearance of favoritism (both `chgnet` and `atomate2` are Materials Project-supported projects), and to avoid dependency conflicts between MLFF packages.
+If you need a sense of which forcefields are compatible, you can use the [pyproject.toml](https://github.com/materialsproject/atomate2/blob/a8bc6505e439503a114f5346aec916aafae7f27b/pyproject.toml#L90) to see which versions are grouped together for testing.
+
 Most of `Maker` classes using the forcefields inherit from `atomate2.forcefields.utils.ForceFieldMixin` to specify which forcefield to use.
 The `ForceFieldMixin` mixin provides the following configurable parameters:
 

pyproject.toml

@@ -54,17 +54,10 @@ defects = [
     "pymatgen-analysis-defects>=2024.5.11",
     "python-ulid>=2.7",
 ]
+
 forcefields = [
     "ase>=3.26.0",
-    "calorine>=3.0",
-    "mace-torch>=0.3.3",
-    "matgl>=1.2.1",
-    "torchdata<=0.7.1", # TODO: remove when issue fixed
-    "quippy-ase>=0.9.14",
-    "mattersim>=1.0.1",
-    "sevenn>=0.9.3",
-    "deepmd-kit>=2.1.4",
-    "upet>=0.2.1",
+    "chgnet>=0.4.2",
 ]
 torchsim = [
     "torch-sim-atomistic==0.5.0; python_version >= '3.12'"
@@ -98,10 +91,12 @@ strict-openff = [
 strict-forcefields-generic = [
     "calorine==3.3; python_version >= '3.12'",
     "calorine==3.1; python_version < '3.12'",
+    "chgnet==0.4.2",
     "quippy-ase==0.10.2",
     "sevenn==0.12.1",
     "deepmd-kit==3.1.3",
-    "tensorflow-cpu==2.21.0",
+    "tensorflow-cpu==2.21.0; sys_platform == 'linux'",
+    "tensorflow==2.21.0; sys_platform == 'darwin' or sys_platform == 'win32'",
     "mattersim==1.2.1",
     "ase<3.28.0", # TODO: remove, required for mattersim because of ase.constraints import
     "wandb>=0.24.0", # required for mattersim
@@ -113,7 +108,7 @@ strict-forcefields-torch-limited = [
     # That enforces a simultaneous pin on torch / torchdata
     # Linux users can acces newer versions of dgl / torch / torchdata via conda.
     # Mac / Windows users will need to install from source
-    "dgl==2.2.1; sys_platform == 'darwin' or sys_platform == 'win32'",
+    "dgl==2.2.0; sys_platform == 'darwin' or sys_platform == 'win32'",
     "dgl<=2.4.0; sys_platform == 'linux'",
     "torch==2.2.2; sys_platform == 'darwin' or sys_platform == 'win32'",
     "torch==2.2.0; sys_platform == 'linux'",

src/atomate2/forcefields/utils.py

@@ -9,6 +9,7 @@
 from functools import cached_property
 from importlib import import_module
 from importlib.metadata import PackageNotFoundError, version
+from importlib.util import find_spec
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -290,6 +291,15 @@ def ase_calculator(
                 calculator = getattr(import_module(_mod), _cls, None)(**kwargs)
 
             case MLFF.CHGNet | MLFF.M3GNet | MLFF.MATPES_R2SCAN | MLFF.MATPES_PBE:
+                if calculator_name == MLFF.CHGNet:
+                    # Legacy interface to `chgnet` package
+                    try:
+                        from chgnet.model.dynamics import CHGNetCalculator
+
+                        return CHGNetCalculator(**kwargs)
+                    except ImportError:
+                        pass
+
                 import matgl
 
                 match calculator_name:
@@ -299,13 +309,12 @@ def ase_calculator(
                     case MLFF.CHGNet:
                         path = kwargs.get("path", "CHGNet-MPtrj-2023.12.1-2.7M-PES")
                         matgl.config.BACKEND = "DGL"
+
                         warnings.warn(
                             "The CHGNet functionality in atomate2 has been migrated "
                             "from the `chgnet` package to `matgl` to ensure continuing "
                             "support. If you want to use the `chgnet` package, "
-                            "`pip install chgnet` and then specify "
-                            '`calculator_meta = {"@module": "chgnet.model.dynamics", '
-                            '"@callable": "CHGNetCalculator"}`',
+                            "`pip install chgnet`",
                             stacklevel=2,
                         )
                     case MLFF.MATPES_R2SCAN | MLFF.MATPES_PBE:
@@ -444,7 +453,13 @@ def _get_pkg_name(calculator_meta: MLFF | dict[str, Any]) -> str | None:
         match calculator_meta:
             case MLFF.Allegro | MLFF.Nequip:
                 ff_pkg = "nequip"
-            case MLFF.CHGNet | MLFF.M3GNet | MLFF.MATPES_PBE | MLFF.MATPES_R2SCAN:
+            case MLFF.CHGNet:
+                # Check if CHGNet is installed
+                try:
+                    ff_pkg = next(pkg for pkg in ("chgnet", "matgl") if find_spec(pkg))
+                except StopIteration:
+                    ff_pkg = None
+            case MLFF.M3GNet | MLFF.MATPES_PBE | MLFF.MATPES_R2SCAN:
                 ff_pkg = "matgl"
             case MLFF.DeepMD:
                 ff_pkg = "deepmd-kit"

tests/forcefields/flows/test_phonon.py

@@ -1,4 +1,5 @@
 import os
+from importlib.util import find_spec
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
@@ -135,6 +136,8 @@ def test_phonon_wf_force_field(
         }
     )
 
+    is_matgl_chgnet = find_spec("matgl") is not None
+
     phonon_kwargs = dict(
         use_symmetrized_structure="conventional",
         create_thermal_displacements=False,
@@ -181,7 +184,9 @@ def test_phonon_wf_force_field(
 
     assert_allclose(
         ph_bs_dos_doc.free_energies,
-        [4440.74345, 4172.361432, 2910.000404, 720.739896, -2194.234779],
+        [4440.74345, 4172.361432, 2910.000404, 720.739896, -2194.234779]
+        if is_matgl_chgnet
+        else [5271.300306, 5162.674841, 4353.717375, 2698.616337, 343.125174],
         atol=1000,
     )
 
@@ -215,7 +220,9 @@ def test_phonon_wf_force_field(
     assert ph_bs_dos_doc.phonopy_settings.kpoint_density_dos == 7_000
     assert_allclose(
         ph_bs_dos_doc.entropies,
-        [0.0, 7.374244, 17.612124, 25.802735, 32.209433],
+        [0.0, 7.374244, 17.612124, 25.802735, 32.209433]
+        if is_matgl_chgnet
+        else [0.0, 3.733666, 12.536534, 20.344558, 26.627292],
         atol=2,
     )
     assert_allclose(

tests/forcefields/test_jobs.py

@@ -1,6 +1,7 @@
 import importlib
 from contextlib import nullcontext
 from importlib.metadata import version as get_imported_version
+from importlib.util import find_spec
 from pathlib import Path
 
 import numpy as np
@@ -38,7 +39,10 @@ def test_maker_initialization(mlff):
 
 
 @pytest.mark.skipif(
-    dgl is None or not mlff_is_installed("CHGNet"),
+    # test to see if CHGNet is installed, or that matgl is installed without dgl
+    # Note that this should be the only test for interface with
+    # the legacy `chgnet` package
+    not mlff_is_installed("CHGNet") or (mlff_is_installed("M3GNet") and dgl is None),
     reason="CHGNet requires DGL which is not installed",
 )
 def test_chgnet_static_maker(si_structure):
@@ -48,17 +52,21 @@ def test_chgnet_static_maker(si_structure):
         ionic_step_data=("structure", "energy"),
     ).make(si_structure)
 
+    pkg_name = "matgl" if find_spec("matgl") else "chgnet"
+
     # run the flow or job and ensure that it finished running successfully
     responses = run_locally(job, ensure_success=True)
 
     # validate job outputs
     output1 = responses[job.uuid][1].output
     assert isinstance(output1, ForceFieldTaskDocument)
-    assert output1.output.energy == approx(-10.7907495, rel=1e-4)
+    assert output1.output.energy == approx(
+        -10.7907495 if pkg_name == "matgl" else -10.6275053, rel=1e-4
+    )
     assert output1.output.ionic_steps[-1].magmoms is None
     assert output1.output.n_steps == 1
 
-    assert output1.forcefield_version == get_imported_version("matgl")
+    assert output1.forcefield_version == get_imported_version(pkg_name)
 
 
 @pytest.mark.skipif(