[fix] (Hopefully) fixed GitHub Action for CI

alexurba · alexurba · commit f845c1049db8 · 2026-03-19T16:07:41.000-04:00
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -11,6 +11,8 @@ permissions:
 env:
   ENV_NAME: aenet-torch
   PYTHON_VERSION: "3.11"
+  TORCH_VERSION: "2.9.0"
+  PYG_WHEEL_URL: "https://data.pyg.org/whl/torch-2.9.0+cpu.html"
 
 jobs:
   unit-tests:
@@ -35,7 +37,10 @@ jobs:
       - name: Install project and dependencies
         run: |
           micromamba run -n "${ENV_NAME}" python -m pip install --upgrade pip
-          micromamba run -n "${ENV_NAME}" python -m pip install -e ".[dev,torch]"
+          micromamba run -n "${ENV_NAME}" python -m pip install -e ".[dev]"
+          micromamba run -n "${ENV_NAME}" python -m pip install "torch==${TORCH_VERSION}"
+          micromamba run -n "${ENV_NAME}" python -m pip install \
+            torch-scatter torch-cluster -f "${PYG_WHEEL_URL}"
 
       - name: Run general pytest suite
         run: |
@@ -63,7 +68,10 @@ jobs:
       - name: Install project and dependencies
         run: |
           micromamba run -n "${ENV_NAME}" python -m pip install --upgrade pip
-          micromamba run -n "${ENV_NAME}" python -m pip install -e ".[dev,torch]"
+          micromamba run -n "${ENV_NAME}" python -m pip install -e ".[dev]"
+          micromamba run -n "${ENV_NAME}" python -m pip install "torch==${TORCH_VERSION}"
+          micromamba run -n "${ENV_NAME}" python -m pip install \
+            torch-scatter torch-cluster -f "${PYG_WHEEL_URL}"
 
       - name: Run pytest docs examples
         run: |
@@ -115,7 +123,10 @@ jobs:
       - name: Install project and dependencies
         run: |
           micromamba run -n "${ENV_NAME}" python -m pip install --upgrade pip
-          micromamba run -n "${ENV_NAME}" python -m pip install -e ".[dev,torch]"
+          micromamba run -n "${ENV_NAME}" python -m pip install -e ".[dev]"
+          micromamba run -n "${ENV_NAME}" python -m pip install "torch==${TORCH_VERSION}"
+          micromamba run -n "${ENV_NAME}" python -m pip install \
+            torch-scatter torch-cluster -f "${PYG_WHEEL_URL}"
 
       - name: Prepare disposable notebook worktree
         run: |
diff --git a/README.md b/README.md
@@ -24,24 +24,44 @@ For development, install in editable mode:
 
 Optional: enable PyTorch-based features (featurization and training)
 
-- From a source checkout:
+Installation matrix:
+
+- Core PyTorch support only
+
+        $ pip install '.[torch]'
+
+  This is sufficient for the core `aenet.torch_training` API, but not for the
+  PyG-backed featurization and neighbor-list stack.
+
+- Full CPU PyTorch stack (recommended for docs examples and notebooks)
 
         $ pip install '.[torch]'
+        $ pip install torch-scatter torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}+cpu.html
 
-  or for development:
+- Full CUDA PyTorch stack
 
-        $ pip install -e '.[dev,torch]'
+        $ pip install '.[torch]'
+        $ pip install torch-scatter torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
+
+- Development install with the full CPU stack
+
+        $ pip install -e '.[dev]'
+        $ pip install '.[torch]'
+        $ pip install torch-scatter torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}+cpu.html
 
 If pip reports build/version conflicts for torch-scatter/torch-cluster:
 - Ensure your installed torch version matches available wheels (CPU/CUDA).
 - You can install prebuilt wheels from the PyG index, for example:
 
         $ pip install torch-scatter torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
 
-  Replace ${TORCH} (e.g., 2.3.1) and ${CUDA} (e.g., cu121 or cpu) with your environment.
+  Replace ${TORCH} (for example, 2.9.0) and ${CUDA} (for example, cu124 or
+  cpu) with your environment.
 
 Behavior without PyTorch extras:
-- aenet.torch_featurize and aenet.torch_training can be imported, but accessing symbols raises a clear ImportError suggesting: pip install 'aenet[torch]'.
+- aenet.torch_featurize and PyG-backed training workflows can be imported, but
+  accessing symbols raises a clear ImportError with guidance to install core
+  torch plus matching `torch-scatter` / `torch-cluster` wheels.
 - AtomicStructure.get_neighbors() automatically falls back to a NumPy neighbor list with a RuntimeWarning.
 
 To use the featurization capabilities, the main ænet package needs to be installed separately as described [elsewhere](http://ann.atomistic.net/documentation/).  And `aenet-python` needs to be made aware of the ænet installation path, for example, using its interactive configuration tool
diff --git a/docs/source/dev/neighbor_lists.rst b/docs/source/dev/neighbor_lists.rst
@@ -16,11 +16,11 @@ and optional GPU execution, see ``notebooks/example-07-neighbor-list.ipynb``.
 .. note::
 
    The features described here make use of PyTorch.  Make sure to
-   install ænet with the ``[torch]`` extra as described in
-   :doc:`/usage/installation`.  For use without PyTorch, ``aenet-python`` also
-   provides a (less efficient) pure-Python neighbor list implementation in
-   ``aenet.geometry.nblist``, which can be used with ``AtomicStructure``
-   objects.
+   install core torch support plus the matching ``torch-scatter`` and
+   ``torch-cluster`` wheels as described in :doc:`/usage/installation`.  For
+   use without PyTorch, ``aenet-python`` also provides a (less efficient)
+   pure-Python neighbor list implementation in ``aenet.geometry.nblist``,
+   which can be used with ``AtomicStructure`` objects.
 
 Quick Start
 -----------
diff --git a/docs/source/usage/installation.rst b/docs/source/usage/installation.rst
@@ -38,6 +38,36 @@ explicitly with
 
    $ pip install ".[torch]"
 
+This installs the core PyTorch dependency only.  The PyTorch-based
+featurization and neighbor-list features also require the PyG extension
+packages ``torch-scatter`` and ``torch-cluster`` matched to the local torch
+build.
+
+Recommended installation matrix:
+
+.. sourcecode:: console
+
+   # Core PyTorch support only
+   $ pip install ".[torch]"
+
+   # Full CPU PyTorch stack
+   $ pip install ".[torch]"
+   $ pip install torch-scatter torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}+cpu.html
+
+   # Full CUDA PyTorch stack
+   $ pip install ".[torch]"
+   $ pip install torch-scatter torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
+
+Replace ``${TORCH}`` with the installed torch version (for example ``2.9.0``)
+and ``${CUDA}`` with the matching CUDA tag (for example ``cu124``).  If you
+are developing from a source checkout, the equivalent editable install is:
+
+.. sourcecode:: console
+
+   $ pip install -e ".[dev]"
+   $ pip install ".[torch]"
+   $ pip install torch-scatter torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}+cpu.html
+
 
 2. Configure ænet Fortran Binaries
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/docs/source/usage/torch_featurization.rst b/docs/source/usage/torch_featurization.rst
@@ -4,7 +4,8 @@ PyTorch-Based Featurization
 .. note::
 
    Featurization as described here makes use of PyTorch.  Make sure to
-   install ænet with the ``[torch]`` extra as described in :doc:`installation`.
+   install core torch support plus the matching ``torch-scatter`` and
+   ``torch-cluster`` wheels as described in :doc:`installation`.
 
 .. note::
 
diff --git a/docs/source/usage/torch_inference.rst b/docs/source/usage/torch_inference.rst
@@ -6,7 +6,9 @@ Using ANN Potentials with the PyTorch Implementation
 .. note::
 
    Inference as described here makes use of PyTorch.  Make sure to
-   install ænet with the ``[torch]`` extra as described in :doc:`installation`.
+   install core torch support as described in :doc:`installation`.  The
+   descriptor-backed workflows shown here also rely on the matching
+   ``torch-scatter`` and ``torch-cluster`` wheels.
 
 .. note::
 
diff --git a/docs/source/usage/torch_training.rst b/docs/source/usage/torch_training.rst
@@ -9,7 +9,9 @@ and automatic differentiation for forces.
 .. note::
 
    Training as described here makes use of PyTorch.  Make sure to
-   install ænet with the ``[torch]`` extra as described in :doc:`installation`.
+   install core torch support as described in :doc:`installation`.  Most
+   descriptor-based training workflows also require the matching
+   ``torch-scatter`` and ``torch-cluster`` wheels.
 
 .. note::
 
diff --git a/notebooks/example-04-torch-featurization.ipynb b/notebooks/example-04-torch-featurization.ipynb
@@ -6,7 +6,7 @@
    "source": [
     "# PyTorch Featurization\n",
     "\n",
-    "**Note: `aenet-python` needs to be installed with the `[torch]` requirements (`pip install aenet[torch]`) for this notebook to work.**\n",
+    "**Note: `aenet-python` needs core torch support (`pip install aenet[torch]`) plus matching `torch-scatter` and `torch-cluster` wheels from `https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html` for this notebook to work.**\n",
     "\n",
     "This notebook demonstrates the longer PyTorch featurization workflows that complement the compact CPU-only examples in `docs/source/usage/torch_featurization.rst`, including file-based input, batch processing, optional GPU execution, and gradient computation."
    ]
diff --git a/notebooks/example-05-torch-training.ipynb b/notebooks/example-05-torch-training.ipynb
@@ -7,7 +7,7 @@
    "source": [
     "# Training with PyTorch\n",
     "\n",
-    "**Note:** `aenet-python` needs to be installed with the `[torch]` requirements (`pip install aenet[torch]`) for this notebook to work.\n",
+    "**Note:** `aenet-python` needs core torch support (`pip install aenet[torch]`) plus matching `torch-scatter` and `torch-cluster` wheels from `https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html` for this notebook to work.\n",
     "\n",
     "This notebook demonstrates end-to-end PyTorch training on the TiO2 example set.\n",
     "It demonstrates file-backed workflows, explicit `CachedStructureDataset` usage, and\n",
diff --git a/notebooks/example-06-torch-inference.ipynb b/notebooks/example-06-torch-inference.ipynb
@@ -7,7 +7,7 @@
    "source": [
     "# Inference with the PyTorch implementation\n",
     "\n",
-    "**Note: `aenet-python` needs to be installed with the `[torch]` requirements (`pip install aenet[torch]`) for this notebook to work.**\n",
+    "**Note: `aenet-python` needs core torch support (`pip install aenet[torch]`) plus matching `torch-scatter` and `torch-cluster` wheels from `https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html` for this notebook to work.**\n",
     "\n",
     "Predicting energies and forces with the PyTorch API is nearly identical to the Fortran/Python API."
    ]
diff --git a/src/aenet/_optional.py b/src/aenet/_optional.py
@@ -5,16 +5,27 @@
 subpackages like `aenet.torch_featurize` and `aenet.torch_training` can be
 imported safely even when the PyTorch stack is not installed.
 
-Users can install extras with:
+Users can install core torch support with:
     pip install "aenet[torch]"
+
+PyG-backed features additionally require `torch-scatter` and
+`torch-cluster` from the matching wheel index at:
+    https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
 """
 
 from __future__ import annotations
 
-from typing import Tuple, Optional
+PYG_WHEEL_INDEX = "https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html"
+TORCH_INSTALL_HINT = "pip install 'aenet[torch]'"
+TORCH_STACK_INSTALL_HINT = (
+    "Install core torch with: "
+    f"{TORCH_INSTALL_HINT}. "
+    "Then install torch-scatter and torch-cluster from the matching "
+    f"PyG wheel index: {PYG_WHEEL_INDEX}"
+)
 
 
-def _safe_import(module: str) -> Tuple[bool, Optional[Exception]]:
+def _safe_import(module: str) -> tuple[bool, Exception | None]:
     try:
         __import__(module)
         return True, None
@@ -37,9 +48,10 @@ def is_torch_cluster_available() -> bool:
     return ok
 
 
-def torch_status() -> Tuple[bool, str]:
+def torch_status() -> tuple[bool, str]:
     """
     Returns (available, reason_if_unavailable).
+
     Only checks core PyTorch.
     """
     ok, err = _safe_import("torch")
@@ -50,9 +62,10 @@ def torch_status() -> Tuple[bool, str]:
     return False, reason
 
 
-def torch_stack_status() -> Tuple[bool, str]:
+def torch_stack_status() -> tuple[bool, str]:
     """
-    Returns (available, reason_if_unavailable) for full featurization stack:
+    Returns (available, reason_if_unavailable) for full featurization stack.
+
       - torch
       - torch-scatter
       - torch-cluster
@@ -75,8 +88,9 @@ def torch_stack_status() -> Tuple[bool, str]:
             False,
             "Missing optional dependencies: "
             + ", ".join(missing)
-            + ". Install with: pip install 'aenet[torch]' "
-              "(ensure torch/CPU/CUDA wheels match your system)",
+            + ". "
+            + TORCH_STACK_INSTALL_HINT
+            + " (ensure torch/CPU/CUDA wheels match your system)",
         )
     return True, ""
 
@@ -86,7 +100,7 @@ def require_torch(feature: str = "this feature") -> None:
     if not ok:
         raise ImportError(
             f"{feature} requires PyTorch. {reason}. "
-            "Install with: pip install 'aenet[torch]'"
+            f"Install with: {TORCH_INSTALL_HINT}"
         )
 
 
@@ -95,5 +109,5 @@ def require_torch_stack(feature: str = "this feature") -> None:
     if not ok:
         raise ImportError(
             f"{feature} requires the PyTorch extras. {reason} "
-            "Install with: pip install 'aenet[torch]'"
+            + TORCH_STACK_INSTALL_HINT
         )
diff --git a/src/aenet/geometry/structure.py b/src/aenet/geometry/structure.py
@@ -1211,7 +1211,10 @@ def get_neighbors(self, i, cutoff, return_self=True, frame=-1):
             warnings.warn(
                 "Falling back to NumPy neighbor list because PyTorch "
                 "extras are unavailable or failed at runtime "
-                f"({e}). Install with: pip install 'aenet[torch]'",
+                f"({e}). Install core torch with: pip install 'aenet[torch]'. "
+                "PyG-backed neighbor lists additionally require "
+                "torch-scatter and torch-cluster from "
+                "https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html",
                 RuntimeWarning,
             )
             from ..nblist.neighborlist import NeighborList  # type: ignore
diff --git a/src/aenet/torch_featurize/__init__.py b/src/aenet/torch_featurize/__init__.py
@@ -1,18 +1,22 @@
 """
 PyTorch-based featurization for aenet.
 
-This subpackage is optional. To use it, install the PyTorch extras:
+This subpackage is optional. To use it, install core torch support:
     pip install "aenet[torch]"
 
+PyG-backed featurization also requires `torch-scatter` and `torch-cluster`
+from the matching wheel index:
+    https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
+
 When the PyTorch stack (torch, torch-scatter, torch-cluster) is not
 installed, importing `aenet.torch_featurize` succeeds, but accessing
 symbols will raise a clear ImportError with installation guidance.
 """
 
 from importlib import import_module
-from typing import Any, Dict, Tuple
+from typing import Any
 
-from .._optional import torch_stack_status
+from .._optional import TORCH_STACK_INSTALL_HINT, torch_stack_status
 
 __all__ = [
     "TorchNeighborList",
@@ -27,7 +31,7 @@
 ]
 
 # Map exported names to (relative_module, attr_name)
-_NAME_TO_SPEC: Dict[str, Tuple[str, str]] = {
+_NAME_TO_SPEC: dict[str, tuple[str, str]] = {
     # neighbor list
     "TorchNeighborList": ("..torch_nblist", "TorchNeighborList"),
     # chebyshev bases
@@ -50,7 +54,7 @@ def __getattr__(name: str) -> Any:  # PEP 562 lazy attribute access
         if not ok:
             raise ImportError(
                 f"{name} requires the PyTorch extras. {reason} "
-                "Install with: pip install 'aenet[torch]'"
+                + TORCH_STACK_INSTALL_HINT
             )
         rel_mod, attr = _NAME_TO_SPEC[name]
         mod = import_module(rel_mod, __name__)
diff --git a/src/aenet/torch_training/__init__.py b/src/aenet/torch_training/__init__.py
@@ -6,11 +6,13 @@
 
 When PyTorch is not installed, importing `aenet.torch_training` succeeds,
 but accessing symbols will raise a clear ImportError with installation
-guidance.
+guidance. Descriptor-based training workflows also require the PyG
+extensions (`torch-scatter` and `torch-cluster`) used by
+`aenet.torch_featurize`.
 """
 
 from importlib import import_module
-from typing import Any, Dict, Tuple
+from typing import Any
 
 from .._optional import require_torch
 
@@ -43,7 +45,7 @@
 ]
 
 # Map exported names to (relative_module, attr_name)
-_NAME_TO_SPEC: Dict[str, Tuple[str, str]] = {
+_NAME_TO_SPEC: dict[str, tuple[str, str]] = {
     # config
     "SGD": (".config", "SGD"),
     "Adam": (".config", "Adam"),

Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,7 @@`
`6`	`6`	`"source": [`
`7`	`7`	`"# PyTorch Featurization\n",`
`8`	`8`	`"\n",`
`9`		- "Note: `aenet-python` needs to be installed with the `[torch]` requirements (`pip install aenet[torch]`) for this notebook to work.\n",
	`9`	+ "Note: `aenet-python` needs core torch support (`pip install aenet[torch]`) plus matching `torch-scatter` and `torch-cluster` wheels from `https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html` for this notebook to work.\n",
`10`	`10`	`"\n",`
`11`	`11`	"This notebook demonstrates the longer PyTorch featurization workflows that complement the compact CPU-only examples in `docs/source/usage/torch_featurization.rst`, including file-based input, batch processing, optional GPU execution, and gradient computation."
`12`	`12`	`]`
Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@`
`7`	`7`	`"source": [`
`8`	`8`	`"# Inference with the PyTorch implementation\n",`
`9`	`9`	`"\n",`
`10`		- "Note: `aenet-python` needs to be installed with the `[torch]` requirements (`pip install aenet[torch]`) for this notebook to work.\n",
	`10`	+ "Note: `aenet-python` needs core torch support (`pip install aenet[torch]`) plus matching `torch-scatter` and `torch-cluster` wheels from `https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html` for this notebook to work.\n",
`11`	`11`	`"\n",`
`12`	`12`	`"Predicting energies and forces with the PyTorch API is nearly identical to the Fortran/Python API."`
`13`	`13`	`]`