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,6 +101,14 @@ 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{-2r}{\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
     """
@@ -96,6 +122,18 @@ def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarra
             np.pi * atomic_gaussian_width_sq
         ) ** (3 / 4)
 
+    def compute_derivative(
+        self, positions: Union[float, np.ndarray]
+    ) -> Union[float, np.ndarray]:
+        atomic_gaussian_width_sq = self.atomic_gaussian_width**2
+        return (
+            -2
+            * positions
+            / atomic_gaussian_width_sq
+            * np.exp(-0.5 * positions**2 / atomic_gaussian_width_sq)
+            / (np.pi * atomic_gaussian_width_sq) ** (3 / 4)
+        )
+
 
 class LodeDensity(AtomicDensityBase):
     r"""Smeared power law density, as used in LODE.
@@ -143,6 +181,7 @@ def __init__(self, atomic_gaussian_width: float, potential_exponent: int):
 
         self.potential_exponent = potential_exponent
         self.atomic_gaussian_width = atomic_gaussian_width
+        self.atomic_gaussian_width_sq = atomic_gaussian_width**2
 
     def _short_range(self, a, x):
         return 1 / a - x / (a + 1) + x**2 / (2 * (a + 2))
@@ -165,3 +204,8 @@ def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarra
                 self._short_range(a, x),
                 self._long_range(a, x),
             )
+
+    def compute_derivative(
+        self, positions: Union[float, np.ndarray]
+    ) -> Union[float, np.ndarray]:
+        raise NotImplementedError("Compute derivative is not implemented yet.")

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
@@ -674,18 +674,62 @@ def integrand(
         )
 
     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(
+                integrand_position: float, u: 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(
+                integrand_position: float, u: 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):