Add radial integral for custom densities

python/rascaline/rascaline/utils/splines/atomic_density.py

@@ -52,12 +52,22 @@ class AtomicDensityBase(ABC):
 
     @abstractmethod
     def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
-        """Compute the atomic density arising from atoms at ``positions``
+        """Compute the atomic density arising from atoms at ``positions``.
 
         :param positions: positions to evaluate the atomic densities
         :returns: evaluated atomic density
         """
 
+    @abstractmethod
+    def compute_derivative(
+        self, positions: Union[float, np.ndarray]
+    ) -> Union[float, np.ndarray]:
+        """Derivative of the atomic density arising from atoms at ``positions``.
+
+        :param positions: positions to evaluate the derivatives atomic densities
+        :returns: evaluated derivative of the atomic density with respect to positions
+        """
+
 
 class DeltaDensity(AtomicDensityBase):
     r"""Delta atomic densities of the form :math:`g(r)=\delta(r)`."""
@@ -67,6 +77,14 @@ def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarra
             "Compute function of the delta density should never called directly."
         )
 
+    def compute_derivative(
+        self, positions: Union[float, np.ndarray]
+    ) -> Union[float, np.ndarray]:
+        raise ValueError(
+            "Compute derivative function of the delta density should never called "
+            "directly."
+        )
+
 
 class GaussianDensity(AtomicDensityBase):
     r"""Gaussian atomic density function.
@@ -83,18 +101,41 @@ class GaussianDensity(AtomicDensityBase):
 
             \|g\|^2 = \int \mathrm{d}^3\boldsymbol{r} |g(r)|^2 = 1\,,
 
+    The derivatives of the Gaussian atomic density with respect to the position is
+
+    .. math::
+
+        g^\prime(r) =
+            \frac{\partial g(r)}{\partial r} = \frac{-r}{\sigma^2(\pi
+            \sigma^2)^{3/4}}e^{-\frac{r^2}{2\sigma^2}} \,.
+
     :param atomic_gaussian_width: Width of the atom-centered gaussian used to create the
         atomic density
     """
 
     def __init__(self, atomic_gaussian_width: float):
         self.atomic_gaussian_width = atomic_gaussian_width
 
-    def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
+    def _compute(
+        self, positions: Union[float, np.ndarray], derivative: bool = False
+    ) -> Union[float, np.ndarray]:
         atomic_gaussian_width_sq = self.atomic_gaussian_width**2
-        return np.exp(-0.5 * positions**2 / atomic_gaussian_width_sq) / (
-            np.pi * atomic_gaussian_width_sq
-        ) ** (3 / 4)
+        x = positions**2 / (2 * atomic_gaussian_width_sq)
+
+        density = np.exp(-x) / (np.pi * atomic_gaussian_width_sq) ** (3 / 4)
+
+        if derivative:
+            density *= -positions / atomic_gaussian_width_sq
+
+        return density
+
+    def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
+        return self._compute(positions=positions, derivative=False)
+
+    def compute_derivative(
+        self, positions: Union[float, np.ndarray]
+    ) -> Union[float, np.ndarray]:
+        return self._compute(positions=positions, derivative=True)
 
 
 class LodeDensity(AtomicDensityBase):
@@ -108,10 +149,11 @@ class LodeDensity(AtomicDensityBase):
                \frac{\gamma\left( \frac{p}{2}, \frac{r^2}{2\sigma^2} \right)}
                     {r^p},
 
-    where :math:`\Gamma(z)` is the Gamma function and :math:`\gamma(a, x)` is the
-    incomplete lower Gamma function. However its evaluation at :math:`r=0` is
-    problematic because :math:`g(r)` is of the form :math:`0/0`. For practical
-    implementations, it is thus more convenient to rewrite the density as
+    where :math:`p` is the potential exponent, :math:`\Gamma(z)` is the Gamma function
+    and :math:`\gamma(a, x)` is the incomplete lower Gamma function. However its
+    evaluation at :math:`r=0` is problematic because :math:`g(r)` is of the form
+    :math:`0/0`. For practical implementations, it is thus more convenient to rewrite
+    the density as
 
     .. math::
 
@@ -123,10 +165,10 @@ class LodeDensity(AtomicDensityBase):
                         & x \geq 10^{-5}
                 \end{cases}
 
-    It is convenient to use the expression for sufficiently small :math:`x` since the
-    relative weight of the first neglected term is on the order of :math:`1/6x^3`.
-    Therefore, the threshold :math:`x = 10^{-5}` leads to relative errors on the order
-    of the machine epsilon.
+    where :math:`a=p/2`. It is convenient to use the expression for sufficiently small
+    :math:`x` since the relative weight of the first neglected term is on the order of
+    :math:`1/6x^3`. Therefore, the threshold :math:`x = 10^{-5}` leads to relative
+    errors on the order of the machine epsilon.
 
     :param atomic_gaussian_width: Width of the atom-centered gaussian used to create the
         atomic density
@@ -141,27 +183,55 @@ def __init__(self, atomic_gaussian_width: float, potential_exponent: int):
         if not HAS_SCIPY:
             raise ValueError("LodeDensity requires scipy to be installed")
 
-        self.potential_exponent = potential_exponent
         self.atomic_gaussian_width = atomic_gaussian_width
+        self.potential_exponent = potential_exponent
 
-    def _short_range(self, a, x):
-        return 1 / a - x / (a + 1) + x**2 / (2 * (a + 2))
+    def _short_range(
+        self, a: float, x: Union[float, np.ndarray], derivative: bool = False
+    ):
+        if derivative:
+            return -1 / (a + 1) + x / (a + 2)
+        else:
+            return 1 / a - x / (a + 1) + x**2 / (2 * (a + 2))
 
-    def _long_range(self, a, x):
-        return gamma(a) * gammainc(a, x) / x**a
+    def _long_range(
+        self, a: float, x: Union[float, np.ndarray], derivative: bool = False
+    ):
+        if derivative:
+            return gamma(a) * gammainc(a, x) / x**a
+        else:
+            return -gamma(a) / x * (np.exp(-x) + a * gammainc(a, x) / x**a)
 
-    def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
+    def _compute(
+        self, positions: Union[float, np.ndarray], derivative: bool = False
+    ) -> Union[float, np.ndarray]:
         if self.potential_exponent == 0:
-            return GaussianDensity.compute(self, positions=positions)
+            return GaussianDensity._compute(
+                self, positions=positions, derivative=derivative
+            )
         else:
             atomic_gaussian_width_sq = self.atomic_gaussian_width**2
             a = self.potential_exponent / 2
             x = positions**2 / (2 * atomic_gaussian_width_sq)
 
             prefac = 1 / gamma(a) / (2 * atomic_gaussian_width_sq) ** a
 
-            return prefac * np.where(
+            density = prefac * np.where(
                 x < 1e-5,
-                self._short_range(a, x),
-                self._long_range(a, x),
+                self._short_range(a, x, derivative=derivative),
+                self._long_range(a, x, derivative=derivative),
             )
+
+            # add inner derivative: ∂x/∂r
+            if derivative:
+                density *= positions / atomic_gaussian_width_sq
+
+            return density
+
+    def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
+        return self._compute(positions=positions, derivative=False)
+
+    def compute_derivative(
+        self, positions: Union[float, np.ndarray]
+    ) -> Union[float, np.ndarray]:
+        return self._compute(positions=positions, derivative=True)

python/rascaline/rascaline/utils/splines/splines.py

@@ -43,8 +43,8 @@
 
 
 try:
-    from scipy.integrate import quad, quad_vec
-    from scipy.special import spherical_in, spherical_jn
+    from scipy.integrate import dblquad, quad, quad_vec
+    from scipy.special import legendre, spherical_in, spherical_jn
 
     HAS_SCIPY = True
 except ImportError:
@@ -593,7 +593,7 @@ def radial_integral_derivative(
         elif type(self.density) is GaussianDensity:
             return self._radial_integral_gaussian_derivative(n, ell, positions)
         else:
-            return self._radial_integral_custom_derivative(n, ell, positions)
+            return self._radial_integral_custom(n, ell, positions, derivative=True)
 
     def _radial_integral_delta(
         self, n: int, ell: int, positions: np.ndarray
@@ -632,9 +632,10 @@ def integrand(
         return (
             prefactor
             * quad_vec(
-                f=lambda x: integrand(x, n, ell, positions),
+                f=integrand,
                 a=0,
                 b=self.basis.integration_radius,
+                args=(n, ell, positions),
             )[0]
         )
 
@@ -666,26 +667,71 @@ def integrand(
         return atomic_gaussian_width_sq**-1 * (
             prefactor
             * quad_vec(
-                f=lambda x: integrand(x, n, ell, positions),
+                f=integrand,
                 a=0,
                 b=self.basis.integration_radius,
+                args=(n, ell, positions),
             )[0]
             - positions * self._radial_integral_gaussian(n, ell, positions)
         )
 
     def _radial_integral_custom(
-        self, n: int, ell: int, positions: np.ndarray, derivative: bool
+        self, n: int, ell: int, positions: np.ndarray, derivative: bool = False
     ) -> np.ndarray:
-        raise NotImplementedError(
-            "Radial integral with custom atomic densities is not implemented yet!"
-        )
 
-    def _radial_integral_custom_derivative(
-        self, n: int, ell: int, positions: np.ndarray, derivative: bool
-    ) -> np.ndarray:
-        raise NotImplementedError(
-            "Radial integral with custom atomic densities is not implemented yet!"
-        )
+        P_ell = legendre(ell)
+
+        if derivative:
+
+            def integrand(
+                u: float, integrand_position: float, n: int, ell: int, position: float
+            ) -> float:
+                arg = np.sqrt(
+                    integrand_position**2
+                    + position**2
+                    - 2 * integrand_position * position * u
+                )
+
+                return (
+                    integrand_position**2
+                    * self.basis.compute(n, ell, integrand_position)
+                    * P_ell(u)
+                    * (position - u * integrand_position)
+                    * self.density.compute_derivative(arg)
+                    / arg
+                )
+
+        else:
+
+            def integrand(
+                u: float, integrand_position: float, n: int, ell: int, position: float
+            ) -> float:
+                arg = np.sqrt(
+                    integrand_position**2
+                    + position**2
+                    - 2 * integrand_position * position * u
+                )
+
+                return (
+                    integrand_position**2
+                    * self.basis.compute(n, ell, integrand_position)
+                    * P_ell(u)
+                    * self.density.compute(arg)
+                )
+
+        radial_integral = np.zeros(len(positions))
+
+        for i, position in enumerate(positions):
+            radial_integral[i], _ = dblquad(
+                func=integrand,
+                a=0,
+                b=self.basis.integration_radius,
+                gfun=-1,
+                hfun=1,
+                args=(n, ell, position),
+            )
+
+        return 2 * np.pi * radial_integral
 
 
 class LodeSpliner(RadialIntegralSplinerBase):
@@ -799,9 +845,10 @@ def integrand(
             )
 
         return quad_vec(
-            f=lambda x: integrand(x, n, ell, positions),
+            f=integrand,
             a=0,
             b=self.basis.integration_radius,
+            args=(n, ell, positions),
         )[0]
 
     def radial_integral_derivative(
@@ -817,9 +864,10 @@ def integrand(
             )
 
         return quad_vec(
-            f=lambda x: integrand(x, n, ell, positions),
+            f=integrand,
             a=0,
             b=self.basis.integration_radius,
+            args=(n, ell, positions),
         )[0]
 
     @property
@@ -848,5 +896,5 @@ def integrand(integrand_position: float, n: int) -> np.ndarray:
             func=integrand,
             a=0,
             b=self.basis.integration_radius,
-            args=(n),
+            args=(n,),
         )[0]

python/rascaline/tests/utils/atomic_density.py

@@ -0,0 +1,18 @@
+import numpy as np
+import pytest
+from numpy.testing import assert_allclose
+
+from rascaline.utils import GaussianDensity
+
+
+pytest.importorskip("scipy")
+
+
+@pytest.mark.parametrize("density", [GaussianDensity(atomic_gaussian_width=1.2)])
+def test_derivative(density):
+    ppositions = np.linspace(0, 5, num=1000)
+    dens = density.compute(ppositions)
+    grad_ana = density.compute_derivative(ppositions)
+    grad_numerical = np.gradient(dens, ppositions)
+
+    assert_allclose(grad_numerical, grad_ana, atol=1e-3)