Update torch (stfc#415)

* Allow updated CHGNet/torch * Update documentation for conflicting extras * Split up extra MLIP tests * Fix and tidy extra tests * Make MACE an extra * Update GitHub Actions * Fix pynvml installation issue * Add missing MACE installation to workflows * Update docs * Remove require torch-dft dependency * Update developer docs * Add skip for alignn download failure * Update GitHub workflows
alinelena · Feb 24, 2025 · 09dff33 · 09dff33
1 parent 55e0ea1
commit 09dff33
Show file tree

Hide file tree

Showing 18 changed files with 598 additions and 378 deletions.
diff --git a/.github/workflows/ci-dev.yml b/.github/workflows/ci-dev.yml
@@ -23,7 +23,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
 
     - name: Install dependencies
-      run: uv sync --all-extras
+      run: uv sync --extra all
 
     # Full reinstall in case the version number is the same as the last release
     - name: Uninstall release versions

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -15,19 +15,30 @@ jobs:
       - uses: actions/checkout@v4
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v4
+        uses: astral-sh/setup-uv@v5
         with:
-          version: "0.5.7"
+          version: "0.6.2"
           python-version: ${{ matrix.python-version }}
 
-      - name: Install dependencies
-        run: uv sync --all-extras
+      - name: Install "all" dependencies
+        run: uv sync --extra all
+
+      - name: Run test suite
+        env:
+          # show timings of tests
+          PYTEST_ADDOPTS: "--durations=0"
+        run: uv run pytest --cov janus_core --cov-append .
+
+      - name: Install dgl dependencies
+        run: |
+          uv sync --extra mace --extra m3gnet --extra alignn
+          uv pip install --reinstall pynvml
 
       - name: Run test suite
         env:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
-        run: uv run pytest --run-extra-mlips --cov janus_core --cov-append .
+        run: uv run pytest tests/test_{mlip_calculators,single_point,eos}.py
 
       - name: Report coverage to Coveralls
         uses: coverallsapp/github-action@v2
@@ -59,9 +70,9 @@ jobs:
                <(sed -nr '/automodule/{s/.*::\s*//;p}' docs/source/apidoc/* | sort)
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v4
+        uses: astral-sh/setup-uv@v5
         with:
-          version: "0.5.7"
+          version: "0.6.2"
           python-version: "3.12"
 
       - name: Install pandoc
@@ -80,9 +91,9 @@ jobs:
       - uses: actions/checkout@v4
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v4
+        uses: astral-sh/setup-uv@v5
         with:
-          version: "0.5.7"
+          version: "0.6.2"
           python-version: "3.12"
 
       - name: Install dependencies

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
@@ -29,9 +29,9 @@ jobs:
         uses: pandoc/actions/setup@v1
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v4
+        uses: astral-sh/setup-uv@v5
         with:
-          version: "0.5.7"
+          version: "0.6.2"
           python-version: "3.12"
 
       - name: Install dependencies

diff --git a/.github/workflows/mac.yml b/.github/workflows/mac.yml
@@ -16,16 +16,27 @@ jobs:
       - uses: actions/checkout@v4
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v4
+        uses: astral-sh/setup-uv@v5
         with:
-          version: "0.5.7"
+          version: "0.6.2"
           python-version: ${{ matrix.python-version }}
 
-      - name: Install dependencies
-        run: uv sync --all-extras
+      - name: Install "all" dependencies
+        run: uv sync --extra all
 
       - name: Run test suite
         env:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
-        run: uv run pytest --run-extra-mlips
+        run: uv run pytest
+
+      - name: Install dgl dependencies
+        run: |
+          uv sync --extra mace --extra m3gnet --extra alignn
+          uv pip install --reinstall pynvml
+
+      - name: Run test suite
+        env:
+          # show timings of tests
+          PYTEST_ADDOPTS: "--durations=0"
+        run: uv run pytest tests/test_{mlip_calculators,single_point,eos}.py
diff --git a/.github/workflows/publish-on-pypi.yml b/.github/workflows/publish-on-pypi.yml
@@ -27,13 +27,13 @@ jobs:
         uses: pandoc/actions/setup@v1
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v4
+        uses: astral-sh/setup-uv@v5
         with:
-          version: "0.5.7"
+          version: "0.6.2"
           python-version: "3.12"
 
       - name: Install dependencies
-        run: uv sync --all-extras
+        run: uv sync --extras all
 
       - name: Build docs with tutorials
         run: cd docs && uv run make all
@@ -65,9 +65,9 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v4
+        uses: astral-sh/setup-uv@v5
         with:
-          version: "0.5.7"
+          version: "0.6.2"
           python-version: "3.12"
 
       - name: Install dependencies

diff --git a/README.md b/README.md
@@ -25,24 +25,8 @@ Tools for machine learnt interatomic potentials
 
 ### Dependencies
 
-`janus-core` dependencies currently include:
-
-- Python >= 3.10
-- ASE >= 3.24
-- mace-torch = 0.3.10
-- chgnet = 0.3.8 (optional)
-- matgl = 1.1.3 (optional)
-- sevenn = 0.10.3 (optional)
-- alignn = 2024.5.27 (optional)
-- nequip = 0.6.1 (optional)
-- deepmd-kit = dpa3-alpha (optional)
-- orb-models = 0.4.2 (optional)
-
 All required and optional dependencies can be found in [pyproject.toml](pyproject.toml).
 
-> [!NOTE]
-> Where possible, we expect to update pinned MLIP dependencies to match their latest releases, subject to any required API fixes.
-
 
 ### Installation
 
@@ -58,15 +42,24 @@ To get all the latest changes, `janus-core` can also be installed from GitHub:
 python3 -m pip install git+https://github.com/stfc/janus-core.git
 ```
 
-By default, MACE is the only MLIP installed.
+By default, no machine learnt interatomic potentials (MLIPs) will be installed with `janus-core`. These can be installed separately, or as `extras`.
 
-Other MLIPs can be installed as `extras`. For example, to install CHGNet and M3GNet, run:
+For example, to install MACE, CHGNet, and SevenNet, run:
 
 ```python
-python3 -m pip install janus-core[chgnet,m3gnet]
+python3 -m pip install janus-core[mace,chgnet,sevennet]
 ```
 
-or to install all supported MLIPs:
+> [!WARNING]
+> `matgl` and `alignn` depend on [dgl](https://github.com/dmlc/dgl?tab=readme-ov-file), which no
+> longer publishes to PyPI. If `janus-core` is installed with either of these extras, PyTorch will
+> automatically be set to 2.2.0 to ensure compatibility. However, this is incompatible with `chgnet`,
+> and may limit the available features in others, including `mace`. To use `matgl` and/or `alignn` with
+> more recent PyTorch release, please refer to the
+> [installation documentation](https://stfc.github.io/janus-core/user_guide/installation.html).
+
+
+To install all MLIPs that do not depend on `dgl`:
 
 ```python
 python3 -m pip install janus-core[all]
@@ -93,11 +86,11 @@ Current and planned features include:
   - MACE
   - M3GNet
   - CHGNet
-  - ALIGNN (experimental)
-  - SevenNet (experimental)
-  - NequIP (experimental)
-  - DPA3 (experimental)
-  - Orb (experimental)
+  - ALIGNN
+  - SevenNet
+  - NequIP
+  - DPA3
+  - Orb
 - [x] Single point calculations
 - [x] Geometry optimisation
 - [x] Molecular Dynamics
@@ -308,7 +301,7 @@ We recommend installing uv for dependency management when developing for `janus-
 ```shell
 git clone https://github.com/stfc/janus-core
 cd janus-core
-uv sync --all-extras # Create a virtual environment and install all dependencies
+uv sync --extras all # Create a virtual environment and install dependencies
 source .venv/bin/activate
 pre-commit install  # Install pre-commit hooks
 pytest -v  # Discover and run all tests

diff --git a/docs/source/developer_guide/get_started.rst b/docs/source/developer_guide/get_started.rst
@@ -27,11 +27,11 @@ Dependencies useful for development can then be installed by running::
 
 Extras, such as optional MLIPs, can also be installed by running::
 
-    uv sync -p 3.12 --extra alignn --extra sevennet
+    uv sync -p 3.12 --extra chgnet --extra sevennet
 
-or to install all supported MLIPs::
+or to install all MLIPs that do not depend on ``dgl``::
 
-    uv sync -p 3.12 --all-extras
+    uv sync -p 3.12 --extras all
 
 
 Using uv
@@ -71,6 +71,11 @@ Packages in the ``dev`` dependency group allow tests to be run locally using ``p
 
     pytest -v
 
+.. note::
+
+    MACE must be installed for tests to run successfully. All other MLIPs are optional.
+
+
 Alternatively, tests can be run in separate virtual environments using ``tox``::
 
     tox run -e ALL