orthonormalize subblocks when multiple species are present (#18)

* orthonormalize subblocks when multiple species are present

* added comments explaining orthonormalizing each subblock per neighbor and refactored the code to create a boolean-array mask variable called neighbor_mask

anisoap/representations/radial_basis.py

@@ -222,20 +222,30 @@ def orthonormalize_basis(self, features: TensorMap):
             )
             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
-            )
+            # Each block's `properties` dimension contains radial channels for each neighbor species
+            # Hence we have to iterate through each neighbor species and orthonormalize the block in subblocks
+            # Each subblock is indexed using the neighbor_mask boolean array.
+            neighbors = np.unique(block.properties["neighbor_species"])
+            for neighbor in neighbors:
+                l = label["angular_channel"]
+                neighbor_mask = block.properties["neighbor_species"] == neighbor
+                n_arr = block.properties["n"][neighbor_mask].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[:, :, neighbor_mask] *= 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[:, :, neighbor_mask] = np.einsum(
+                    "ijk,kl->ijl",
+                    block.values[:, :, neighbor_mask],
+                    orthonormalization_matrix,
+                )
 
         return features