Solid harm prefact (#38)

* added the missing solid harmonic prefactor

* added some comments on why dividing through by this prefactor is necessary

* added test to ensure that the expansion coefficients correctly lead the gaussian to converge. also added more documentation to moment_generator, clarifying the normalization constant

* added metatensor to requirements.txt

* corrected the testing code. rtol should not be np.inf otherwise the test never fails. Also moved the scaling outside of the inner loop, significantly speeds up the iterations

* Update anisoap/representations/ellipsoidal_density_projection.py

from a for loop to an np divide.

Co-authored-by: Rose K. Cersonsky <[email protected]>

* minor formatting

* Revert "Update anisoap/representations/ellipsoidal_density_projection.py
"
Turns out, can't use numpy.divide here, because each entry in sph_to_cart has a different shape that numpy can't figure out how to resolve.
This reverts commit e896cbf.

* linter

---------

Co-authored-by: Rose K. Cersonsky <[email protected]>

anisoap/representations/ellipsoidal_density_projection.py

@@ -79,6 +79,14 @@ def pairwise_ellip_expansion(
     keys = [tuple(i) + (l,) for i in keys for l in range(lmax + 1)]
     num_ns = radial_basis.get_num_radial_functions()
     maxdeg = np.max(np.arange(lmax + 1) + 2 * np.array(num_ns))
+
+    # This prefactor is the solid harmonics prefactor, that we need to divide by later.
+    # This is needed because spherical_to_cartesian calculates solid harmonics Rlm = sqrt((4pi)/(2l+1)) * r^l*Ylm
+    # Our expansion coefficients from the inner product does not have this prefactor included, so we divide it later.
+    solid_harm_prefact = np.sqrt((4 * np.pi) / (np.arange(lmax + 1) * 2 + 1))
+    scaled_sph_to_cart = []
+    for l in range(lmax + 1):
+        scaled_sph_to_cart.append(sph_to_cart[l] / solid_harm_prefact[l])
     for center_types in types:
         for neighbor_types in types:
             if (center_types, neighbor_types) in neighbor_list.keys:
@@ -129,7 +137,11 @@ def pairwise_ellip_expansion(
                         deg = l + 2 * (num_ns[l] - 1)
                         moments_l = moments[: deg + 1, : deg + 1, : deg + 1]
                         values_ldict[l].append(
-                            np.einsum("mnpqr, pqr->mn", sph_to_cart[l], moments_l)
+                            np.einsum(
+                                "mnpqr, pqr->mn",
+                                scaled_sph_to_cart[l],
+                                moments_l,
+                            )
                         )
 
                 for l in tqdm(

anisoap/utils/moment_generator.py

@@ -129,8 +129,11 @@ def compute_moments_inefficient_implementation(A, a, maxdeg):
         <x^{n_0} * y^{n_1} * z^{n_2}> = \int(x^{n_0} * y^{n_1} * z^{n_2}) * \exp(-0.5*(r-a).T@cov@(r-a)) dxdydz
         \sum_{i=1}^{\\infty} x_{i}
 
-        Note that the term "moments" in probability theory are defined for normalized Gaussian distributions.
-        Here, we take the Gaussian
+        Note that the term "moments" in probability theory are defined for normalized Gaussian distributions. These
+        recursive relations find the moments of a normalized Gaussian, but we actually want to find the moments of
+        the unnormalized underlying gaussian (as seen in the equation above) because calculating expansion
+        coefficients of any density should not involve normalization. Therefore, we scale the end-result by
+        global_factor, which is the reciprocal of the normalizing constant in front of any gaussian.
     """
     # Make sure that the provided arrays have the correct dimensions and properties
     assert A.shape == (3, 3), "Dilation matrix needs to be 3x3"

anisoap/utils/spherical_to_cartesian.py

@@ -64,6 +64,17 @@ def binom(n, k):
 
 
 def spherical_to_cartesian(lmax, num_ns):
+    """
+    Finds the coefficients for the cartesian polynomial form of solid harmonics R_{lm} = sqrt((4pi)/(2l+1))*r^l*Y_{lm}.
+    Note that our AniSOAP expansion does not contain the sqrt((4pi)/(2l+1)), so in calculating expansion coefficients,
+    we need to divide by that coefficient.
+    Args:
+        lmax:
+        num_ns:
+
+    Returns:
+
+    """
     assert len(num_ns) == lmax + 1
 
     # Initialize array in which to store all

pyproject.toml

@@ -14,12 +14,10 @@ commands =
 
 [tool.coverage.run]
 branch = true
-data_file = 'tests/.coverage'
+data_file = "tests/.coverage"
 
 [tool.coverage.report]
-include = [
-    "anisoap/*"
-]
+include = ["anisoap/*"]
 
 [tool.coverage.xml]
-output = 'tests/coverage.xml'
+output = "tests/coverage.xml"

tests/requirements.txt

@@ -4,3 +4,4 @@ pytest
 scipy
 tqdm
 git+https://github.com/Luthaf/rascaline.git
+metatensor

tests/test_convergence.py

@@ -0,0 +1,100 @@
+import numpy as np
+from ase import Atoms
+import pytest
+from numpy.testing import assert_allclose
+from anisoap.representations import EllipsoidalDensityProjection
+import metatensor
+
+from scipy.special import sph_harm
+
+
+class TestGaussianConvergence:
+    sigmas = np.array([1.0, 2.0, 3.0])
+
+    @pytest.mark.parametrize("sigma", sigmas)
+    def test_single_atom_isotropic_convergence(self, sigma):
+        """
+        Test that coefficients are correct, such that the approximation using these coefficients in the expansion
+        can reasonably recreate the original atomic gaussian (in this case, an unnormalized gaussian with sigma=1)
+        """
+
+        frame = Atoms(positions=np.array([[0, 0, 0]]), numbers=[0])
+        frame.arrays["quaternions"] = np.array([[1, 0, 0, 0]])
+
+        frame.arrays[r"c_diameter[1]"] = 2 * sigma * np.ones(1)
+        frame.arrays[r"c_diameter[2]"] = 2 * sigma * np.ones(1)
+        frame.arrays[r"c_diameter[3]"] = 2 * sigma * np.ones(1)
+        frames = [frame]
+
+        # An rgw that's a little bigger than the atom gaussian sigma converges well and is sufficient for this test.
+        rgw = sigma + 0.5
+
+        max_radials = range(10)
+        errs = []
+        for max_radial in max_radials:
+            HYPER_PARAMETERS = {
+                "max_angular": 1,
+                "max_radial": max_radial,
+                "radial_basis_name": "gto",
+                "rotation_type": "quaternion",
+                "rotation_key": "quaternions",
+                "cutoff_radius": 1.0,
+                "radial_gaussian_width": rgw,
+                "basis_rcond": 1e-14,
+                "basis_tol": 1e-1,
+            }
+            representation = EllipsoidalDensityProjection(**HYPER_PARAMETERS)
+            descriptor_raw = representation.transform(frames, normalize=True)
+            descriptor = metatensor.operations.sum_over_samples(
+                descriptor_raw, sample_names="center"
+            )
+
+            def evaluate_bases(r):
+                def real_sphharm(m, l, theta, phi):
+                    if m < 0:
+                        return (
+                            np.sqrt(2)
+                            * (-1.0) ** m
+                            * np.imag(sph_harm(np.abs(m), l, theta, phi))
+                        )
+                    elif m == 0:
+                        # this is real anyway, just doing this so we don't get an annoying "discarding imaginary warning"
+                        return np.real(sph_harm(m, l, theta, phi))
+                    else:
+                        return (
+                            np.sqrt(2)
+                            * (-1.0) ** m
+                            * np.real(sph_harm(m, l, theta, phi))
+                        )
+
+                bases = descriptor.copy()
+                for key, block in bases.items():
+                    l = key["angular_channel"]
+                    for m in block.components[0][
+                        "spherical_component_m"
+                    ]:  # same as range(-l, l+1)
+                        # evaluated at rhat = (theta, phi) = (0, 0) -- ok if gaussian is isotropic.
+                        ylm = real_sphharm(m, l, 0, 0)
+                        for n in block.properties["n"]:
+                            rnl = representation.radial_basis.get_basis(r).flatten()[n]
+                            block.values[0, m, n] = ylm * rnl
+                return bases
+
+        # Now, we perform our reconstruction
+        approx = []
+        r_mesh = np.linspace(0, 5, 100)
+        actual = np.exp(-(r_mesh**2) / (2 * sigma**2))
+
+        for r in r_mesh:
+            # For now, have to do an inefficient pointwise evaluation, unfortunately.
+            bases = evaluate_bases(np.array([r]))
+            approx.append(np.sum(bases.block(0).values * descriptor.block(0).values))
+        approx = np.array(approx)
+        errs.append(np.sum((approx - actual) ** 2) / len(approx))
+
+        # assert that errors are monotonically decreasing. As we increase the number of terms in the expansion,
+        # we should get a better and better approximation.
+        assert np.all(errs[:-1] >= errs[1:])
+
+        # as long as each element differs by a (somewhat) small absolute amount, we're good.
+        assert_allclose(approx, actual, atol=1e-03)