Merged changes from the most recent 'main'

.github/workflows/lint.yml

@@ -34,3 +34,4 @@ jobs:
       - name: Check imports
         run: |
           isort anisoap/*/*py -m 3 --tc --fgw --up -e -l 88 --check
+          isort anisoap/*py -m 3 --tc --fgw --up -e -l 88 --check

.github/workflows/tests.yml

@@ -27,6 +27,9 @@ jobs:
       - name: Run tests
         run: |
           pytest -v tests/.
-      - uses: codecov/codecov-action@v1
+      - uses: codecov/codecov-action@v3
         with:
           file: ./tests/coverage.xml
+        env:
+          name: CODECOV_TOKEN
+          value: ${{ secrets.CODECOV_TOKEN }}

README.md

@@ -1,5 +1,11 @@
-# anisoap
-A Python Package for Computing the Smooth Overlap of Anisotropic Positions
+AniSOAP
+=======
+
+<a href="https://github.com/cersonsky-lab/anisoap/actions?query=workflow%3ATest">
+  <img src="https://github.com/cersonsky-lab/anisoap/workflows/Test/badge.svg"/>
+</a><a href="https://codecov.io/gh/cersonsky-lab/anisoap/">
+  <img src="https://codecov.io/gh/cersonsky-lab/anisoap/branch/main/graph/badge.svg?token=UZJPJG34SM" />
+</a>
 
 ## Installation
 
@@ -30,3 +36,11 @@ Please run pytest and check that all tests pass before pushing new changes to th
 
     pytest tests/.
 
+Contributors
+------------
+
+Thanks goes to all people that make AniSOAP possible:
+
+<a href="https://github.com/cersonsky-lab/anisoap/graphs/contributors">
+  <img src="https://contrib.rocks/image?repo=cersonsky-lab/anisoap" />
+</a>

anisoap/__init__.py

@@ -1,3 +1,6 @@
-from anisoap import representations, utils, lib
+from anisoap import (
+    representations,
+    utils
+)
 
 __version__ = "0.0.0"

anisoap/representations/radial_basis.py

@@ -1,4 +1,90 @@
+import warnings
+
 import numpy as np
+import scipy.linalg
+from equistore.core import TensorMap
+from scipy.special import gamma
+
+
+def inverse_matrix_sqrt(matrix: np.array):
+    """
+    Returns the inverse matrix square root.
+    The inverse square root of the overlap matrix (or slices of the overlap matrix) yields the
+    orthonormalization matrix
+    Args:
+        matrix: np.array
+            Symmetric square matrix to find the inverse square root of
+
+    Returns:
+        inverse_sqrt_matrix: S^{-1/2}
+
+    """
+    if not np.allclose(matrix, matrix.conjugate().T):
+        raise ValueError("Matrix is not hermitian")
+    eva, eve = np.linalg.eigh(matrix)
+
+    if (eva < 0).any():
+        raise ValueError(
+            "Matrix is not positive semidefinite. Check that a valid gram matrix is passed."
+        )
+    return eve @ np.diag(1 / np.sqrt(eva)) @ eve.T
+
+
+def gto_square_norm(n, sigma):
+    """
+    Compute the square norm of GTOs (inner product of itself over R^3).
+    An unnormalized GTO of order n is \phi_n = r^n * e^{-r^2/(2*\sigma^2)}
+    The square norm of the unnormalized GTO has an analytic solution:
+    <\phi_n | \phi_n> = \int_0^\infty dr r^2 |\phi_n|^2 = 1/2 * \sigma^{2n+3} * \Gamma(n+3/2)
+    Args:
+        n: order of the GTO
+        sigma: width of the GTO
+
+    Returns:
+        square norm: The square norm of the unnormalized GTO
+    """
+    return 0.5 * sigma ** (2 * n + 3) * gamma(n + 1.5)
+
+
+def gto_prefactor(n, sigma):
+    """
+    Computes the normalization prefactor of an unnormalized GTO.
+    This prefactor is simply 1/sqrt(square_norm_area).
+    Scaling a GTO by this prefactor will ensure that the GTO has square norm equal to 1.
+    Args:
+        n: order of GTO
+        sigma: width of GTO
+
+    Returns:
+        N: normalization constant
+
+    """
+    return np.sqrt(1 / gto_square_norm(n, sigma))
+
+
+def gto_overlap(n, m, sigma_n, sigma_m):
+    """
+    Compute overlap of two *normalized* GTOs
+    Note that the overlap of two GTOs can be modeled as the square norm of one GTO, with an effective
+    n and sigma. All we need to do is to calculate those effective parameters, then compute the normalization.
+    <\phi_n, \phi_m> = \int_0^\infty dr r^2 r^n * e^{-r^2/(2*\sigma_n^2) * r^m * e^{-r^2/(2*\sigma_m^2)
+    = \int_0^\infty dr r^2 |r^{(n+m)/2} * e^{-r^2/4 * (1/\sigma_n^2 + 1/\sigma_m^2)}|^2
+    = \int_0^\infty dr r^2 r^n_{eff} * e^{-r^2/(2*\sigma_{eff}^2)
+    prefactor.
+    ---Arguments---
+    n: order of the first GTO
+    m: order of the second GTO
+    sigma_n: sigma parameter of the first GTO
+    sigma_m: sigma parameter of the second GTO
+
+    ---Returns---
+    S: overlap of the two normalized GTOs
+    """
+    N_n = gto_prefactor(n, sigma_n)
+    N_m = gto_prefactor(m, sigma_m)
+    n_eff = (n + m) / 2
+    sigma_eff = np.sqrt(2 * sigma_n**2 * sigma_m**2 / (sigma_n**2 + sigma_m**2))
+    return N_n * N_m * gto_square_norm(n_eff, sigma_eff)
 
 
 class RadialBasis:
@@ -7,9 +93,9 @@ class RadialBasis:
     This helps to keep a cleaner main code by avoiding if-else clauses
     related to the radial basis.
 
-    TODO: In the long run, this class would precompute quantities like
-    the normalization factors or orthonormalization matrix for the
-    radial basis.
+    Code relating to GTO orthonormalization is heavily inspired by work done in librascal, specifically this
+    codebase here: https://github.com/lab-cosmo/librascal/blob/8405cbdc0b5c72a5f0b0c93593100dde348bb95f/bindings/rascal/utils/radial_basis.py
+
     """
 
     def __init__(self, radial_basis, max_angular, **hypers):
@@ -27,6 +113,12 @@ def __init__(self, radial_basis, max_angular, **hypers):
             num_n = (max_angular - l) // 2 + 1
             self.num_radial_functions.append(num_n)
 
+        # As part of the initialization, compute the orthonormalization matrix for GTOs
+        # If we are using the monomial basis, set self.overlap_matrix equal to None
+        self.overlap_matrix = None
+        if self.radial_basis == "gto":
+            self.overlap_matrix = self.calc_gto_overlap_matrix()
+
     # Get number of radial functions
     def get_num_radial_functions(self):
         return self.num_radial_functions
@@ -57,3 +149,73 @@ def compute_gaussian_parameters(self, r_ij, lengths, rotation_matrix):
             center -= 1 / sigma**2 * np.linalg.solve(precision, r_ij)
 
         return precision, center
+
+    def calc_gto_overlap_matrix(self):
+        """
+        Computes the overlap matrix for GTOs.
+        The overlap matrix is a Gram matrix whose entries are the overlap: S_{ij} = \int_0^\infty dr r^2 phi_i phi_j
+        The overlap has an analytic solution (see above functions).
+        The overlap matrix is the first step to generating an orthonormal basis set of functions (Lodwin Symmetric
+        Orthonormalization). The actual orthonormalization matrix cannot be fully precomputed because each tensor
+        block use a different set of GTOs. Hence, we precompute the full overlap matrix of dim l_max, and while
+        orthonormalizing each tensor block, we generate the respective orthonormal matrices from slices of the full
+        overlap matrix.
+
+        Returns:
+            S: 2D array. The overlap matrix
+        """
+        # Consequence of the floor divide used to compute self.num_radial_functions
+        max_deg = self.max_angular + 1
+        n_grid = np.arange(max_deg)
+        sigma = self.hypers["radial_gaussian_width"]
+        sigma_grid = np.ones(max_deg) * sigma
+        S = gto_overlap(
+            n_grid[:, np.newaxis],
+            n_grid[np.newaxis, :],
+            sigma_grid[:, np.newaxis],
+            sigma_grid[np.newaxis, :],
+        )
+        return S
+
+    def orthonormalize_basis(self, features: TensorMap):
+        """
+        Apply an in-place orthonormalization on the features, using Lodwin Symmetric Orthonormalization.
+        Each block in the features TensorMap uses a GTO set of l + 2n, so we must take the appropriate slices of
+        the overlap matrix to compute the orthonormalization matrix.
+        An instructive example of Lodwin Symmetric Orthonormalization of a 2-element basis set is found here:
+        https://booksite.elsevier.com/9780444594365/downloads/16755_10030.pdf
+
+        Parameters:
+            features: A TensorMap whose blocks' values we wish to orthonormalize. Note that features is modified in place, so a
+            copy of features must be made before the function if you wish to retain the unnormalized values.
+            radial_basis: An instance of RadialBasis
+
+        Returns:
+            normalized_features: features containing values multiplied by proper normalization factors.
+        """
+        # In-place modification.
+        radial_basis_name = self.radial_basis
+        if radial_basis_name != "gto":
+            warnings.warn(
+                f"Normalization has not been implemented for the {radial_basis_name} basis, and features will not be normalized.",
+                UserWarning,
+            )
+            return features
+        for label, block in features.items():
+            l = label["angular_channel"]
+            n_arr = block.properties["n"].flatten()
+            l_2n_arr = l + 2 * n_arr
+            # normalize all the GTOs by the appropriate prefactor first, since the overlap matrix is in terms of
+            # normalized GTOs
+            prefactor_arr = gto_prefactor(
+                l_2n_arr, self.hypers["radial_gaussian_width"]
+            )
+            block.values[:, :, :] = block.values[:, :, :] * prefactor_arr
+
+            gto_overlap_matrix_slice = self.overlap_matrix[l_2n_arr, :][:, l_2n_arr]
+            orthonormalization_matrix = inverse_matrix_sqrt(gto_overlap_matrix_slice)
+            block.values[:, :, :] = np.einsum(
+                "ijk,kl->ijl", block.values, orthonormalization_matrix
+            )
+
+        return features