Equivariant powerspectrum by pair

python/featomic/featomic/clebsch_gordan/__init__.py

@@ -3,4 +3,7 @@
 from ._coefficients import calculate_cg_coefficients  # noqa: F401
 from ._density_correlations import DensityCorrelations  # noqa: F401
 from ._equivariant_power_spectrum import EquivariantPowerSpectrum  # noqa: F401
+from ._equivariant_power_spectrum_by_pair import (  # noqa: F401
+    EquivariantPowerSpectrumByPair,  # noqa: F401
+)  # noqa: F401
 from ._power_spectrum import PowerSpectrum  # noqa: F401

python/featomic/featomic/clebsch_gordan/_coefficients.py

@@ -701,7 +701,7 @@ def _cg_tensor_product_dense(
         # output shape is [(samples q p), lambda]
         output = _cg_couple_dense(tensor_product, o3_lambda, cg_coefficients)
         #  => [samples, (q p), lambda]
-        output = output.reshape(n_s, (n_p * n_q), -1)
+        output = output.reshape(n_s, (n_p * n_q), 2 * o3_lambda + 1)
         #  => [samples, lambda, (q p)]
         output = _dispatch.swapaxes(output, 1, 2)
         result.append(output)