Refactor core components, update CI/infrastructure, and improve test stability

.github/workflows/lint.yml

@@ -4,10 +4,12 @@ on:
   push:
     branches:
     - main
+    - develop
 
   pull_request:
     branches:
     - main
+    - develop
 
 jobs:
   lint:

.github/workflows/release.yml

@@ -14,7 +14,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v4
       with:
-        python-version: '3.11'
+        python-version: '3.13'
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip

.github/workflows/tests.yml

@@ -4,23 +4,25 @@ on:
   push:
     branches:
     - main
+    - develop
 
   pull_request:
     branches:
     - main
+    - develop
 
 jobs:
   build:
 
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.11']
+        python-version: ['3.13']
 
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
         cache: pip
@@ -43,4 +45,8 @@ jobs:
     - name: Upload to coveralls
       if: github.event_name == 'push'
       run: |
-        COVERALLS_REPO_TOKEN=${{ secrets.COVERALLS_TOKEN }} coveralls
+        if [ -n "${{ secrets.COVERALLS_TOKEN }}" ]; then
+          COVERALLS_REPO_TOKEN="${{ secrets.COVERALLS_TOKEN }}" coveralls
+        else
+          echo "Skipping coveralls upload: COVERALLS_TOKEN not set."
+        fi

.pre-commit-config.yaml

@@ -0,0 +1,19 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v6.0.0
+    hooks:
+      - id: end-of-file-fixer
+      - id: mixed-line-ending
+      - id: trailing-whitespace
+
+  - repo: https://github.com/psf/black
+    rev: 26.3.1
+    hooks:
+      - id: black
+        exclude: (.*)/migrations
+
+  - repo: https://github.com/pycqa/flake8
+    rev: 7.3.0
+    hooks:
+      - id: flake8
+        additional_dependencies: ['flake8-pyproject']

.readthedocs.yaml

@@ -20,4 +20,4 @@ python:
    install:
    - requirements: docs/requirements.txt
    - method: pip
-     path: .
+     path: .

README.md

@@ -7,7 +7,7 @@ import e3nn_jax as e3nn
 # Create a random array made of a scalar (0e) and a vector (1o)
 array = e3nn.normal("0e + 1o", jax.random.PRNGKey(0))
 
-print(array)  
+print(array)
 # 1x0e+1x1o [ 1.8160863  -0.75488514  0.33988908 -0.53483534]
 
 # Compute the norms
@@ -59,33 +59,33 @@ The main difference is the presence of the class [`IrrepsArray`](https://e3nn-ja
 - Euclidean Neural Networks
 ```
 @misc{thomas2018tensorfieldnetworksrotation,
-      title={Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds}, 
+      title={Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds},
       author={Nathaniel Thomas and Tess Smidt and Steven Kearnes and Lusann Yang and Li Li and Kai Kohlhoff and Patrick Riley},
       year={2018},
       eprint={1802.08219},
       archivePrefix={arXiv},
       primaryClass={cs.LG},
-      url={https://arxiv.org/abs/1802.08219}, 
+      url={https://arxiv.org/abs/1802.08219},
 }
 
 @misc{weiler20183dsteerablecnnslearning,
-      title={3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data}, 
+      title={3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data},
       author={Maurice Weiler and Mario Geiger and Max Welling and Wouter Boomsma and Taco Cohen},
       year={2018},
       eprint={1807.02547},
       archivePrefix={arXiv},
       primaryClass={cs.LG},
-      url={https://arxiv.org/abs/1807.02547}, 
+      url={https://arxiv.org/abs/1807.02547},
 }
 
 @misc{kondor2018clebschgordannetsfullyfourier,
-      title={Clebsch-Gordan Nets: a Fully Fourier Space Spherical Convolutional Neural Network}, 
+      title={Clebsch-Gordan Nets: a Fully Fourier Space Spherical Convolutional Neural Network},
       author={Risi Kondor and Zhen Lin and Shubhendu Trivedi},
       year={2018},
       eprint={1806.09231},
       archivePrefix={arXiv},
       primaryClass={stat.ML},
-      url={https://arxiv.org/abs/1806.09231}, 
+      url={https://arxiv.org/abs/1806.09231},
 }
 ```
 - e3nn

docs/requirements.txt

@@ -13,4 +13,4 @@ jraph
 nequip_jax @ git+https://github.com/mariogeiger/nequip-jax@1.1.0
 flax
 dm-haiku
-equinox
+equinox

docs/tuto/index.rst

@@ -4,4 +4,4 @@ Tutorial
 .. toctree::
     :maxdepth: 1
 
-    nequip
+    nequip

e3nn_jax/_src/irreps.py

@@ -2,7 +2,7 @@
 import dataclasses
 import itertools
 import math
-from typing import Callable, List, NamedTuple, Optional, Tuple, Union
+from typing import Callable, Union
 
 import jax
 import jax.numpy as jnp
@@ -18,7 +18,8 @@
 
 from .J import Jd
 
-IntoIrrep = Union[int, "Irrep", "MulIrrep", Tuple[int, int]]
+IntoIrrep = Union[int, "Irrep", "MulIrrep", tuple[int, int]]
+SortResult = collections.namedtuple("SortResult", ["irreps", "p", "inv"])
 
 
 @dataclasses.dataclass(init=False, frozen=True)
@@ -350,14 +351,7 @@ def __lt__(self, other):
     MulIrrep,
     str,
     "Irreps",
-    List[
-        Union[
-            str,
-            Irrep,
-            MulIrrep,
-            Tuple[int, IntoIrrep],
-        ]
-    ],
+    list[str | Irrep | MulIrrep | tuple[int, IntoIrrep]],
 ]
 
 
@@ -405,7 +399,7 @@ def __new__(cls, irreps: IntoIrreps = None):
         if isinstance(irreps, Irreps):
             return super().__new__(cls, irreps)
 
-        out: List[MulIrrep] = []
+        out: list[MulIrrep] = []
         if isinstance(irreps, Irrep):
             out.append(MulIrrep(1, Irrep(irreps)))
         elif irreps is None:
@@ -472,7 +466,7 @@ def spherical_harmonics(lmax, p=-1):
         """
         return Irreps([(1, (l, p**l)) for l in range(lmax + 1)])
 
-    def slices(self) -> List[slice]:
+    def slices(self) -> list[slice]:
         r"""List of slices corresponding to indices for each irrep.
 
         Examples:
@@ -649,9 +643,7 @@ def simplify(self) -> "Irreps":
         """
         return self.remove_zero_multiplicities().unify()
 
-    def sort(
-        self,
-    ) -> NamedTuple("Sort", irreps="Irreps", p=Tuple[int, ...], inv=Tuple[int, ...]):
+    def sort(self) -> SortResult:
         r"""Sort the representations.
 
         Returns:
@@ -669,13 +661,12 @@ def sort(
             >>> Irreps("2o + 1e + 0e + 1e").sort().inv
             (2, 1, 3, 0)
         """
-        Ret = collections.namedtuple("sort", ["irreps", "p", "inv"])
         out = [(ir, i, mul) for i, (mul, ir) in enumerate(self)]
         out = sorted(out)
         inv = tuple(i for _, i, _ in out)
         p = perm.inverse(inv)
         irreps = Irreps([(mul, ir) for ir, _, mul in out])
-        return Ret(irreps, p, inv)
+        return SortResult(irreps, p, inv)
 
     def regroup(self) -> "Irreps":
         r"""Regroup the same irreps together.
@@ -702,9 +693,9 @@ def set_mul(self, mul: int) -> "Irreps":
 
     def filter(
         self,
-        keep: Union["Irreps", List[Irrep], Callable[[MulIrrep], bool]] = None,
+        keep: Union["Irreps", list[Irrep], Callable[[MulIrrep], bool]] = None,
         *,
-        drop: Union["Irreps", List[Irrep], Callable[[MulIrrep], bool]] = None,
+        drop: Union["Irreps", list[Irrep], Callable[[MulIrrep], bool]] = None,
         lmax: int = None,
     ) -> "Irreps":
         r"""Filter the irreps.
@@ -839,7 +830,7 @@ def mul_gcd(self) -> int:
         return math.gcd(*[mul for mul, _ in self])
 
     @property
-    def ls(self) -> List[int]:
+    def ls(self) -> list[int]:
         """List of the l values.
 
         Examples:
@@ -1021,10 +1012,7 @@ def __getitem__(self, index: slice) -> Irreps:
 
 
 def _wigner_D_from_angles(
-    l: int,
-    alpha: Optional[jax.Array],
-    beta: Optional[jax.Array],
-    gamma: Optional[jax.Array],
+    l: int, alpha: jax.Array | None, beta: jax.Array | None, gamma: jax.Array | None
 ) -> jax.Array:
     r"""The Wigner-D matrix of the real irreducible representations of :math:`SO(3)`.