From 507168a09a9f5606d3beab195ee0769a06843cf5 Mon Sep 17 00:00:00 2001 From: Philip Loche Date: Fri, 22 Sep 2023 19:59:07 +0200 Subject: [PATCH] Added real and k space spliner base classes --- python/rascaline/rascaline/utils/splines.py | 354 +++++++++++++++++++- 1 file changed, 350 insertions(+), 4 deletions(-) diff --git a/python/rascaline/rascaline/utils/splines.py b/python/rascaline/rascaline/utils/splines.py index af8534b5e..036a04a2f 100644 --- a/python/rascaline/rascaline/utils/splines.py +++ b/python/rascaline/rascaline/utils/splines.py @@ -21,6 +21,15 @@ import numpy as np +try: + from scipy.integrate import quad, quad_vec + from scipy.special import spherical_in, spherical_jn + + HAS_SCIPY = True +except ImportError: + HAS_SCIPY = False + + class RadialIntegralSplinerBase(ABC): """Base class for splining arbitrary radial integrals. @@ -31,6 +40,8 @@ class RadialIntegralSplinerBase(ABC): :parameter max_radial: number of angular components :parameter spline_cutoff: cutoff radius for the spline interpolation. This is also the maximal value that can be interpolated. + :parameter orthonormalization_cutoff: Provide value if the radial integral should be + orthonormalized. If :py:obj:`None` no orthonormalization will be performed. :parameter accuracy: accuracy of the numerical integration and the splining. Accuracy is reached when either the mean absolute error or the mean relative error gets below the ``accuracy`` threshold. @@ -41,13 +52,18 @@ def __init__( max_radial: int, max_angular: int, spline_cutoff: float, + orthonormalization_cutoff: Optional[float], accuracy: float, ): self.max_radial = max_radial self.max_angular = max_angular self.spline_cutoff = spline_cutoff + self.orthonormalization_cutoff = orthonormalization_cutoff self.accuracy = accuracy + if self.orthonormalization_cutoff is not None and not HAS_SCIPY: + raise ValueError("Orthornomalization requires scipy!") + @abstractmethod def _radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: """Method calculating the radial integral.""" @@ -55,7 +71,11 @@ def _radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarra @property def _center_contribution(self) -> Union[None, np.ndarray]: - r"""Contribution of the central atom required for LODE calculations.""" + r"""Contribution of the central atom. + + Required for LODE calculations. The central atom contribution will not be + orthornomalized! + """ return None @@ -64,7 +84,7 @@ def _radial_integral_derivative( ) -> np.ndarray: """Method calculating the derivatice of the radial integral.""" displacement = 1e-6 - mean_abs_positions = np.abs(positions).mean() + mean_abs_positions = np.mean(np.abs(positions)) if mean_abs_positions <= 1.0: raise ValueError( @@ -82,9 +102,74 @@ def _radial_integral_derivative( return (radial_integral_pos - radial_integral_neg) / displacement + def _orthonormalization_matrix( + self, + func: Callable[[int, int, np.ndarray], np.ndarray], + ) -> np.ndarray: + """Orthornomalize a function.""" + # Gram matrix (also called overlap matrix or inner product matrix) + gram_matrix = np.zeros((self.max_angular + 1, self.max_radial, self.max_radial)) + + # Get the normalization constants from the diagonal entries + self.normalizations = np.zeros((self.max_angular + 1, self.max_radial)) + + for ell in range(self.max_angular + 1): + for n1 in range(self.max_radial): + for n2 in range(self.max_radial): + + def integrand( + integrand_positions: np.ndarray, + func: Callable[[int, int, np.ndarray], np.ndarray], + n1: int, + n2: int, + ell: int, + ) -> np.ndarray: + return ( + integrand_positions**2 + * func(n1, ell, integrand_positions) + * func(n2, ell, integrand_positions) + ) + + gram_matrix[ell, n1, n2] = quad( + func=integrand, + a=0, + b=self.orthonormalization_cutoff, + args=( + func, + n1, + n2, + ell, + ), + epsabs=self.accuracy, + epsrel=self.accuracy, + ) + + for ell in range(self.max_angular + 1): + for n in range(self.max_radial): + self.normalizations[ell, n] = 1 / np.sqrt(gram_matrix[ell, n, n]) + + # Rescale orthonormalization matrix to be defined + # in terms of the normalized (but not yet orthonormalized) + # basis functions + gram_matrix[ell, n, :] *= self.normalizations[ell, n] + gram_matrix[ell, :, n] *= self.normalizations[ell, n] + + # Compute orthonormalization matrix + orthonormalization_matrix = np.zeros( + (self.max_angular + 1, self.max_radial, self.max_radial) + ) + for ell in range(self.max_angular + 1): + eigvals, eigvecs = np.linalg.eigh(gram_matrix[ell]) + orthonormalization_matrix[ell] = ( + eigvecs @ np.diag(np.sqrt(1.0 / eigvals)) @ eigvecs.T + ) + + return orthonormalization_matrix + def _value_evaluator_3D( self, positions: np.ndarray, + orthonormalization_matrix: Optional[np.ndarray], derivative: bool, ): values = np.zeros([len(positions), self.max_angular + 1, self.max_radial]) @@ -97,6 +182,15 @@ def _value_evaluator_3D( else: values[:, ell, n] = self._radial_integral(n, ell, positions) + if orthonormalization_matrix is not None: + # For each l channel we do a dot product of the orthonormalization_matrix of + # shape (n, n) with the values which should have the shape (n, n_positions). + # To achieve the correct broadcasting we have to transpose twice. + for ell in range(self.max_angular + 1): + values[:, ell, :] = ( + orthonormalization_matrix[ell] @ values[:, ell, :].T + ).T + return values def compute( @@ -111,11 +205,22 @@ def compute( rascaline calculator. """ + if self.orthonormalization_cutoff is not None: + orthonormalization_matrix = self._orthonormalization_matrix( + self._radial_integral + ) + else: + orthonormalization_matrix = None + def value_evaluator_3D(positions): - return self._value_evaluator_3D(positions, derivative=False) + return self._value_evaluator_3D( + positions, orthonormalization_matrix, derivative=False + ) def derivative_evaluator_3D(positions): - return self._value_evaluator_3D(positions, derivative=True) + return self._value_evaluator_3D( + positions, orthonormalization_matrix, derivative=True + ) if n_spline_points is not None: positions = np.linspace(0, self.spline_cutoff, n_spline_points) @@ -304,6 +409,8 @@ class RadialIntegralFromFunction(RadialIntegralSplinerBase): the maximal value that can be interpolated. :parameter max_radial: number of angular componentss :parameter max_angular: number of radial components + :parameter orthonormalization_cutoff: Provide value if the radial integral should be + orthonormalized. If :py:obj:`None` no orthonormalization will be performed. :parameter radial_integral_derivative: The derivative of the radial integral taking the same paramaters as ``radial_integral``. If it is :py:obj:`None` (default), finite differences are used to calculate the derivative of the radial integral. @@ -368,6 +475,7 @@ def __init__( spline_cutoff: float, max_radial: int, max_angular: int, + orthonormalization_cutoff: Optional[float] = None, radial_integral_derivative: Optional[ Callable[[int, int, np.ndarray], np.ndarray] ] = None, @@ -382,6 +490,7 @@ def __init__( max_radial=max_radial, max_angular=max_angular, spline_cutoff=spline_cutoff, + orthonormalization_cutoff=orthonormalization_cutoff, accuracy=accuracy, ) @@ -401,3 +510,240 @@ def _radial_integral_derivative( return super()._radial_integral_derivative(n, ell, positions) else: return self.radial_integral_derivative_funcion(n, ell, positions) + + +class RealSpaceRadialIntegralSplinerBase(RadialIntegralSplinerBase): + def __init__( + self, + max_radial: int, + max_angular: int, + atomic_gussian_width: Optional[float], + spline_cutoff: float, + orthonormalization_cutoff: Optional[float] = None, + accuracy: float = 1e-8, + ): + if not HAS_SCIPY: + raise ValueError("Spliner class requires scipy!") + + self.atomic_gussian_width = atomic_gussian_width + + if self.orthonormalization_cutoff is None: + self.radial_integral_cutoff = np.inf + else: + self.radial_integral_cutoff = self.orthonormalization_cutoff + + super().__init__( + max_radial=max_radial, + max_angular=max_angular, + spline_cutoff=spline_cutoff, + orthonormalization_cutoff=orthonormalization_cutoff, + accuracy=accuracy, + ) + + @abstractmethod + def _radial_basis( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + ... + + @abstractmethod + def _radial_basis_derivative( + self, n: float, ell: float, integrand_positions: np.ndarray + ) -> np.ndarray: + """Derivative of the radial basis. + + Only required for the delta atomic density.""" + displacement = 1e-6 + mean_abs_positions = np.abs(integrand_positions).mean() + + if mean_abs_positions <= 1.0: + raise ValueError( + "Numerically derivative of the radial integral can not be performed " + "since positions are too small. Mean of the absolute positions is " + f"{mean_abs_positions:.1e} but should be at least 1." + ) + + radial_basis_pos = self._radial_basis( + n, ell, integrand_positions + displacement / 2 + ) + radial_basis_neg = self._radial_basis( + n, ell, integrand_positions - displacement / 2 + ) + + return (radial_basis_pos - radial_basis_neg) / displacement + + def _atomic_density(self, integrand_positions: np.ndarray) -> np.ndarray: + ... + + def _radial_integral_delta( + self, n: int, ell: int, positions: np.ndarray, derivative: bool + ) -> np.ndarray: + if derivative: + return self._radial_basis_derivative(n, ell, positions) + else: + return self._radial_basis(n, ell, positions) + + def _radial_integral_gaussian( + self, n: int, ell: int, positions: np.ndarray, derivative: bool + ) -> np.ndarray: + prefac = ( + 4 + * np.pi + / (np.pi * self.atomic_gaussian_width**2) ** (3 / 4) + * np.exp(-(positions**2) / (2 * self.atomic_gaussian_width**2)) + ) + + def integrand( + integrand_position: float, n: int, ell: int, positions: np.array + ) -> np.ndarray: + return ( + integrand_position**2 + * self._radial_basis(n, ell, integrand_position) + * np.exp( + -(integrand_position**2) / (2 * self.atomic_gaussian_width**2) + ) + * spherical_in( + ell, + integrand_position * positions / self.atomic_gaussian_width**2, + ) + ) + + Inl = quad_vec( + f=integrand, + a=0, + b=self.radial_integral_cutoff, + args=(n, ell, positions), + epsabs=self.accuracy, + epsrel=self.accuracy, + ) + + if not derivative: + return prefac * Inl + else: + + def dintegrand( + integrand_position: float, n: int, ell: int, positions: np.array + ) -> np.ndarray: + return ( + integrand_position**3 + * self._radial_basis(n, ell, integrand_position) + * np.exp( + -(integrand_position**2) + / (2 * self.atomic_gaussian_width**2) + ) + * spherical_in( + ell, + integrand_position + * positions + / self.atomic_gaussian_width**2, + derivative=True, + ) + ) + + # TODO: Check equation + return self.atomic_gaussian_width**-2 * ( + prefac + * quad_vec( + f=dintegrand, + a=0, + b=self.radial_integral_cutoff, + args=(n, ell, positions), + epsabs=self.accuracy, + epsrel=self.accuracy, + ) + - positions * Inl + ) + + def _radial_integral_custom( + self, n: int, ell: int, positions: np.ndarray, derivative: bool + ) -> np.ndarray: + # Here we will need the `_atomic_density` method to perform the integral! + raise NotImplementedError( + "Radial integral with custom atomic densities is not implemented yet!" + ) + + def _radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: + if self.atomic_gussian_width == 0: + return self._radial_integral_delta(n, ell, positions, derivative=False) + elif self.atomic_gussian_width is None: + return self._radial_integral_custom(n, ell, positions, derivative=False) + else: + return self._radial_integral_gaussian(n, ell, positions, derivative=False) + + def _radial_integral_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + if self.atomic_gussian_width == 0: + return self._radial_integral_delta(n, ell, positions, derivative=True) + elif self.atomic_gussian_width is None: + return self._radial_integral_custom(n, ell, positions, derivative=True) + else: + return self._radial_integral_gaussian(n, ell, positions, derivative=True) + + +class FourierSpaceRadialIntegralSplinerBase(RadialIntegralSplinerBase): + def __init__( + self, + max_radial: int, + max_angular: int, + spline_cutoff: float, + orthonormalization_cutoff: Optional[float] = None, + accuracy: float = 1e-8, + ): + if not HAS_SCIPY: + raise ValueError("Spliner class requires scipy!") + + if self.orthonormalization_cutoff is None: + self.radial_integral_cutoff = np.inf + else: + self.radial_integral_cutoff = self.orthonormalization_cutoff + + super().__init__( + max_radial=max_radial, + max_angular=max_angular, + spline_cutoff=spline_cutoff, + orthonormalization_cutoff=orthonormalization_cutoff, + accuracy=accuracy, + ) + + @abstractmethod + def _radial_basis(self, n: float, ell: float, integrand_position: float) -> float: + ... + + def _radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: + def integrand( + integrand_position: float, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + return ( + integrand_position**2 + * self._radial_basis(n, ell, integrand_position) + * spherical_jn(ell, integrand_position * positions) + ) + + return quad_vec( + f=integrand, + a=0, + b=self.radial_integral_cutoff, + epsabs=self.accuracy, + epsrel=self.accuracy, + ) + + def _radial_integral_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + def integrand( + integrand_position: float, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + return ( + integrand_position**3 + * self._radial_basis(n, ell, integrand_position) + * spherical_jn(ell, integrand_position * positions, derivative=True) + ) + + return quad_vec( + f=integrand, + a=0, + b=self.radial_integral_cutoff, + epsabs=self.accuracy, + epsrel=self.accuracy, + )