From a5f2f9a7c10506a8fd3ee5941a868a24d052364e Mon Sep 17 00:00:00 2001 From: Arthur Lin Date: Tue, 15 Aug 2023 12:22:11 -0500 Subject: [PATCH 1/2] orthonormalize subblocks when multiple species are present --- anisoap/representations/radial_basis.py | 45 ++++++++++++++++--------- 1 file changed, 30 insertions(+), 15 deletions(-) diff --git a/anisoap/representations/radial_basis.py b/anisoap/representations/radial_basis.py index 9c1e208..3b6def0 100644 --- a/anisoap/representations/radial_basis.py +++ b/anisoap/representations/radial_basis.py @@ -202,20 +202,35 @@ 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 - ) + neighbors = np.unique(block.properties["neighbor_species"]) + for neighbor in neighbors: + l = label["angular_channel"] + n_arr = block.properties["n"][ + block.properties["neighbor_species"] == neighbor + ].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.properties["neighbor_species"] == neighbor] = ( + block.values[:, :, block.properties["neighbor_species"] == neighbor] + * 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[ + :, :, block.properties["neighbor_species"] == neighbor + ] = np.einsum( + "ijk,kl->ijl", + block.values[ + :, :, block.properties["neighbor_species"] == neighbor + ], + orthonormalization_matrix, + ) return features From 0e138c67f5c58af6819b348126a967721ed435b0 Mon Sep 17 00:00:00 2001 From: Arthur Lin Date: Wed, 16 Aug 2023 13:12:36 -0500 Subject: [PATCH 2/2] 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 | 21 ++++++++------------- 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/anisoap/representations/radial_basis.py b/anisoap/representations/radial_basis.py index 3b6def0..3c91f78 100644 --- a/anisoap/representations/radial_basis.py +++ b/anisoap/representations/radial_basis.py @@ -202,34 +202,29 @@ def orthonormalize_basis(self, features: TensorMap): ) return features for label, block in features.items(): + # 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"] - n_arr = block.properties["n"][ - block.properties["neighbor_species"] == neighbor - ].flatten() + 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[:, :, block.properties["neighbor_species"] == neighbor] = ( - block.values[:, :, block.properties["neighbor_species"] == neighbor] - * prefactor_arr - ) + 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[ - :, :, block.properties["neighbor_species"] == neighbor - ] = np.einsum( + block.values[:, :, neighbor_mask] = np.einsum( "ijk,kl->ijl", - block.values[ - :, :, block.properties["neighbor_species"] == neighbor - ], + block.values[:, :, neighbor_mask], orthonormalization_matrix, )