Tutorials

.gitignore

@@ -12,3 +12,4 @@ dev_notes/
 .DS_Store
 evals/.venv-torch/
 evals/results/
+.venv_tutorials

docs/guide/tensor_products.md

@@ -34,6 +34,28 @@ output = tp(left_features, edge_attributes)
 Use {class}`~e3nn_mlx.o3.ElementwiseTensorProduct` for aligned channels and
 {class}`~e3nn_mlx.o3.TensorSquare` when both inputs are the same value.
 
+## Symmetry-reduced products
+
+{class}`~e3nn_mlx.o3.ReducedTensorProducts` constructs an orthonormal
+change-of-basis tensor subject to index permutation symmetries such as
+`"ij=ji"` or `"ijk=jik=ikj"`. Its `filter_ir_mid` argument restricts the
+irreps allowed at every sequential coupling stage; `filter_ir_out` restricts
+the retained final irreps. Intermediate filtering prunes contraction paths and
+is not equivalent to slicing an already-constructed output.
+
+```python
+lmax = 4
+bispectrum = o3.ReducedTensorProducts(
+    "ijk=jik=ikj",
+    i=o3.Irreps.spherical_harmonics(lmax),
+    filter_ir_mid=list(o3.Irrep.iterator(lmax=lmax)),
+    filter_ir_out=list(o3.Irrep.iterator(lmax=0)),
+)
+```
+
+This construction retains scalar symmetric triple contractions suitable for a
+bispectrum. 
+
 ## General instructions
 
 The general {class}`~e3nn_mlx.o3.TensorProduct` accepts an explicit output

e3nn_mlx/ops_reduce_tensor.py

@@ -65,13 +65,19 @@ def _irrep_embeddings(irreps: Irreps, np) -> list[tuple[Irrep, Any]]:
     return embeddings
 
 
-def _couple_channels(factors: tuple[Irreps, ...], np) -> list[_CoupledChannel]:
+def _couple_channels(
+    factors: tuple[Irreps, ...],
+    np,
+    filter_ir_mid: frozenset[Irrep] | None = None,
+) -> list[_CoupledChannel]:
     channels = [_CoupledChannel(irrep, embedding) for irrep, embedding in _irrep_embeddings(factors[0], np)]
     for factor in factors[1:]:
         next_channels = []
         for channel in channels:
             for right_irrep, right_basis in _irrep_embeddings(factor, np):
                 for output_irrep in channel.irrep * right_irrep:
+                    if filter_ir_mid is not None and output_irrep not in filter_ir_mid:
+                        continue
                     coefficient = np.asarray(
                         clebsch_gordan(channel.irrep, right_irrep, output_irrep), dtype=np.float64
                     )
@@ -95,14 +101,28 @@ def _nullspace(matrix, np, tolerance: float = 1e-9):
 def _build_change_of_basis(
     formula: str,
     factor_specs: tuple[str, ...],
+    filter_ir_mid_specs: tuple[str, ...] | None = None,
+    filter_ir_out_specs: tuple[str, ...] | None = None,
 ) -> tuple[Irreps, Any]:
     import numpy as np
 
     base, constraints = _parse_formula(formula)
     factors = tuple(Irreps(spec) for spec in factor_specs)
-    channels = _couple_channels(factors, np)
+    filter_ir_mid = (
+        None
+        if filter_ir_mid_specs is None
+        else frozenset(Irrep.parse(spec) for spec in filter_ir_mid_specs)
+    )
+    filter_ir_out = (
+        None
+        if filter_ir_out_specs is None
+        else frozenset(Irrep.parse(spec) for spec in filter_ir_out_specs)
+    )
+    channels = _couple_channels(factors, np, filter_ir_mid)
     grouped: dict[Irrep, list[Any]] = {}
     for channel in channels:
+        if filter_ir_out is not None and channel.irrep not in filter_ir_out:
+            continue
         grouped.setdefault(channel.irrep, []).append(channel.basis)
 
     output_parts = []
@@ -134,15 +154,94 @@ def _build_change_of_basis(
 
 
 class ReducedTensorProducts(mlx_module_base()):
-    """Reduce a tensor product according to a permutation-symmetry formula."""
+    """Reduce a tensor product according to permutation symmetries.
+
+    Parameters
+    ----------
+    formula
+        Index formula describing the tensor symmetry. For example, ``"ij=ji"``
+        selects symmetric rank-two tensors and ``"ijk=jik=ikj"`` selects fully
+        symmetric rank-three tensors.
+    filter_ir_out
+        Optional iterable of allowed final irreps. Channels with other output
+        irreps are omitted from :attr:`irreps_out` and
+        :attr:`change_of_basis`.
+    filter_ir_mid
+        Optional iterable of irreps allowed at every sequential coupling stage,
+        including the final coupling. This prunes contraction paths before the
+        permutation-symmetry reduction and can substantially reduce basis
+        construction work.
+    **irreps
+        Irreps associated with formula indices. If only one index is supplied,
+        its irreps are used for every index in the formula.
+
+    Notes
+    -----
+    ``filter_ir_mid`` is a path constraint, not merely an output slice. It can
+    change which multiplicity channels survive the symmetry reduction.
+
+    Examples
+    --------
+    Construct scalar bispectrum contractions from spherical coefficients:
+
+    >>> from e3nn_mlx import o3
+    >>> lmax = 4
+    >>> bispectrum = o3.ReducedTensorProducts(
+    ...     "ijk=jik=ikj",
+    ...     i=o3.Irreps.spherical_harmonics(lmax),
+    ...     filter_ir_mid=list(o3.Irrep.iterator(lmax=lmax)),
+    ...     filter_ir_out=list(o3.Irrep.iterator(lmax=0)),
+    ... )
+    >>> all(part.ir.l == 0 for part in bispectrum.irreps_out)
+    True
+    """
 
-    def __init__(self, formula: str, **irreps: Irreps | str) -> None:
+    def __init__(
+        self,
+        formula: str,
+        filter_ir_out: list[Irrep | str] | None = None,
+        filter_ir_mid: list[Irrep | str] | None = None,
+        **irreps: Irreps | str,
+    ) -> None:
         super().__init__()
         base, _ = _parse_formula(formula)
         self.formula = formula
         self._labels = base
         self.irreps_in = _factor_irreps(base, irreps)
-        self.irreps_out, data = _build_change_of_basis(formula, tuple(str(value) for value in self.irreps_in))
+        try:
+            self.filter_ir_out = (
+                None
+                if filter_ir_out is None
+                else tuple(Irrep.parse(value) for value in filter_ir_out)
+            )
+        except (TypeError, ValueError) as error:
+            raise ValueError(
+                f"filter_ir_out (={filter_ir_out}) must be an iterable of Irrep values"
+            ) from error
+        try:
+            self.filter_ir_mid = (
+                None
+                if filter_ir_mid is None
+                else tuple(Irrep.parse(value) for value in filter_ir_mid)
+            )
+        except (TypeError, ValueError) as error:
+            raise ValueError(
+                f"filter_ir_mid (={filter_ir_mid}) must be an iterable of Irrep values"
+            ) from error
+        self.irreps_out, data = _build_change_of_basis(
+            formula,
+            tuple(str(value) for value in self.irreps_in),
+            (
+                None
+                if self.filter_ir_mid is None
+                else tuple(str(value) for value in self.filter_ir_mid)
+            ),
+            (
+                None
+                if self.filter_ir_out is None
+                else tuple(str(value) for value in self.filter_ir_out)
+            ),
+        )
         mx, _ = require_mlx()
         self._change_of_basis = mx.array(data, dtype=mx.float32)
 

tests/test_documentation_contract.py

@@ -69,6 +69,8 @@
         "tests/test_upstream_o3_tensor_product.py::test_upstream_tensor_product_unshared_weight_broadcast_and_validation",
     "tensor-product wrappers match explicit contractions":
         "tests/test_tensor_product_correctness.py::test_full_tensor_product_matches_explicit_tensor_product",
+    "reduced-tensor intermediate and output filters":
+        "tests/test_upstream_o3_reduce_tensor.py::test_upstream_reduced_tensor_supports_intermediate_and_output_filters",
     "lazy package import without MLX":
         "tests/test_nn_import_policy.py::test_nn_namespace_imports_without_site_packages",
     "gate-points exact documented configuration":
@@ -176,6 +178,11 @@
         "`e3nn_core` remains importable without MLX.",
         "tests/test_nn_import_policy.py::test_nn_namespace_imports_without_site_packages",
     ),
+    (
+        "docs/guide/tutorial_compatibility.md",
+        "`filter_ir_mid` restricts every sequential Clebsch--Gordan coupling, including the final coupling.",
+        "tests/test_upstream_o3_reduce_tensor.py::test_upstream_reduced_tensor_supports_intermediate_and_output_filters",
+    ),
     (
         "docs/HIGH_LEVEL_API.md",
         "High-level modules preserve the caller's style",