Improvements to CG tensor product

[Luthaf]

The main change here is a separation between the calculation of tensor products and the projection of spherical harmonics with CG coefficients. This will allow using the projection step (cg_couple functions) in separate places, and should be faster by not re-computing tensor product multiple times (as an anecdotal benchmark, density correlation tests went from 13s to 11s on my laptop).

At the same time, this cleans the code around CG tensor product to try to make them easier to follow and with less complex types.

---

📚 Documentation preview 📚: https://rascaline--314.org.readthedocs.build/en/314/

[github-actions[bot]]

Here is a pre-built version of the code in this pull request: wheels.zip, you can install it locally by unzipping wheels.zip and using pip to install the file matching your system

[Luthaf]

This PR also removed the MOPS backend for CG tensor products since it was broken (it segfaults while computing). We will re-enable it later, once mops is stable enough that we can use it on CI.

For reference, here is the latest version of the code before it got removed:

def _cg_tensor_product_mops(
    array_1: Array,
    array_2: Array,
    o3_lambdas: List[int],
    cg_coefficients: TensorMap,
) -> List[Array]:
    """
    Compute the Clebsch-Gordan tensor product of 2 arrays using MOPS' sparse
    accumulation of products.
    """
    import mops

    # Samples dimensions must be the same
    assert array_1.shape[0] == array_2.shape[0]
    assert isinstance(array_1, np.ndarray)
    assert isinstance(array_2, np.ndarray)

    # Infer l1 and l2 from the len of the length of axis 1 of each tensor
    l1 = (array_1.shape[1] - 1) // 2
    l2 = (array_2.shape[1] - 1) // 2

    # Define other useful dimensions
    n_s = array_1.shape[0]  # number of samples
    n_p = array_1.shape[2]  # number of properties in array_1
    n_q = array_2.shape[2]  # number of properties in array_2

    # MOPS sparse accumulation requires some reshaping of the input arrays. See
    # https://github.com/lab-cosmo/mops . Currently only supported for a numpy array
    # backend.
    array_1 = np.repeat(array_1[:, :, :, None], n_q, axis=3).reshape(
        n_s, 2 * l1 + 1, n_p * n_q
    )
    array_2 = np.repeat(array_2[:, :, None, :], n_p, axis=2).reshape(
        n_s, 2 * l2 + 1, n_p * n_q
    )

    array_1 = _dispatch.swapaxes(array_1, 1, 2).reshape(n_s * n_p * n_q, 2 * l1 + 1)
    array_2 = _dispatch.swapaxes(array_2, 1, 2).reshape(n_s * n_p * n_q, 2 * l2 + 1)

    results = []
    for o3_lambda in o3_lambdas:
        # We also need to pass SAP the CG coefficients and m1, m2, and mu indices as 1D
        # arrays. Extract these from the corresponding TensorBlock in the TensorMap CG
        # cache.
        block = cg_coefficients.block({"l1": l1, "l2": l2, "lambda": o3_lambda})
        samples = block.samples

        m1_arr = []
        m2_arr = []
        mu_arr = []
        C_arr = []
        for sample_i, (m1, m2, mu) in enumerate(samples):

            m1_arr.append(int(m1))
            m2_arr.append(int(m2))
            mu_arr.append(int(mu))
            C_arr.append(float(block.values[sample_i, 0]))

        output = mops.sparse_accumulation_of_products(
            A=array_1,
            B=array_2,
            C=np.array(C_arr),
            indices_A=np.array(m1_arr),
            indices_B=np.array(m2_arr),
            indices_output=np.array(mu_arr),
            output_size=2 * o3_lambda + 1,
        )
        assert output.shape == (n_s * n_p * n_q, 2 * o3_lambda + 1)

        # Reshape back
        output = output.reshape(n_s, n_p * n_q, 2 * o3_lambda + 1)
        output = _dispatch.swapaxes(output, 1, 2)

        results.append(output)

    return results

@pytest.mark.skipif(not HAS_MOPS, reason="mops is not installed")
def test_correlate_density_mops_python_sparse_agree():
    """
    Tests for agreement between nu=2 tensors built using both "python-sparse"
    and "mops" CG backend.
    """
    frames = h2o_periodic()
    density = spherical_expansion_small(frames)

    correlation_order = 2
    corr_calculator_python = DensityCorrelations(
        max_angular=SPHEX_HYPERS_SMALL["max_angular"] * correlation_order,
        correlation_order=correlation_order,
        cg_backend="python-sparse",
    )
    corr_calculator_mops = DensityCorrelations(
        max_angular=SPHEX_HYPERS_SMALL["max_angular"] * correlation_order,
        correlation_order=correlation_order,
        cg_backend="mops",
    )
    # NOTE: testing the private function here so we can control the use of
    # sparse v dense CG cache
    n_body_python = corr_calculator_python.compute(density)
    n_body_mops = corr_calculator_mops.compute(density)

    assert metatensor.allclose(n_body_python, n_body_mops, atol=1e-8, rtol=1e-8)

[jwa7]

Also for reference later (maybe useful in docs), the removed docstrings from a couple of functions:

def sparse_combine(
    array_1: Array,
    array_2: Array,
    o3_lambda: int,
    cg_coefficients: TensorMap,
    cg_backend: str,
) -> Array:
    """
    Performs a Clebsch-Gordan combination step on 2 arrays using sparse operations. The
    angular channel of each block is inferred from the size of its component axis, and
    the blocks are combined to the desired output angular channel `o3_lambda` using the
    appropriate Clebsch-Gordan coefficients.
    :param array_1: array with the m values for l1 with shape [n_samples, 2 * l1 + 1,
        n_q_properties]
    :param array_2: array with the m values for l2 with shape [n_samples, 2 * l2 + 1,
        n_p_properties]
    :param o3_lambda: int value of the resulting coupled channel
    :param cg_coefficients: a :py:class:`TensorMap` containing CG coefficients in a
        format for either sparse or dense CG tensor products, as returned by
        :py:func:`calculate_cg_coefficients`. See the function docstring for details on
        the data structure. Only used if ``cg_backend`` is not ``"metadata"``.
    :param cg_backend: specifies the combine backend with sparse CG coefficients. It can
        have the values "python-dense", "python-sparse", "mops" and "metadata". If
        "python-dense" or "python-sparse" is chosen, a dense or sparse combination
        (respectively) of the arrays is performed using either numpy or torch, depending
        on the backend. If "mops" is chosen, a sparse combination of the arrays is
        performed if the external package MOPS is installed. If "metadata" is chosen, no
        combination is perfomed, and an empty array of the correct shape is returned.
    :returns: array of shape [n_samples, (2*o3_lambda+1), q_properties * p_properties]
    """

def dense_combine(
    array_1: Array,
    array_2: Array,
    o3_lambda: int,
    cg_coefficients: TensorMap,
) -> Array:
    """
    Performs a Clebsch-Gordan combination step on 2 arrays using a dense operation. The
    angular channel of each block is inferred from the size of its component axis, and
    the blocks are combined to the desired output angular channel `o3_lambda` using the
    appropriate Clebsch-Gordan coefficients.
    :param array_1: array with the m values for l1 with shape [n_samples, 2 * l1 + 1,
        n_q_properties]
    :param array_2: array with the m values for l2 with shape [n_samples, 2 * l2 + 1,
        n_p_properties]
    :param o3_lambda: int value of the resulting coupled channel
    :param cg_coefficients: a :py:class:`TensorMap` containing CG coefficients in a
        format for either sparse or dense CG tensor products, as returned by
        :py:func:`calculate_cg_coefficients`. See the function docstring for details on
        the data structure. Only used if ``cg_backend`` is not ``"metadata"``.
    :returns: array of shape [n_samples, (2 * o3_lambda + 1), q_properties *
        p_properties]
    """