diff --git a/anisoap/reference/projection_coefficients.py b/anisoap/reference/projection_coefficients.py index 48c6f8f..45ac783 100644 --- a/anisoap/reference/projection_coefficients.py +++ b/anisoap/reference/projection_coefficients.py @@ -1,18 +1,14 @@ import numpy as np -from anisoap.utils.spherical_to_cartesian import spherical_to_cartesian +from ..representations.radial_basis import RadialBasis +from ..utils import quaternion_to_rotation_matrix # missing? +from ..utils import compute_moments_inefficient_implementation +from ..utils.spherical_to_cartesian import spherical_to_cartesian try: from tqdm import tqdm except ImportError: - tqdm = lambda i, **kwargs: i - -from ..utils import ( - compute_moments_diagonal_inefficient_implementation, - compute_moments_inefficient_implementation, - quaternion_to_rotation_matrix, -) -from .radial_basis import RadialBasis + tqdm = lambda i, **_: i class DensityProjectionCalculator: @@ -55,7 +51,7 @@ def __init__( if compute_gradients: raise NotImplementedError("Sorry! Gradients have not yet been implemented") - # Precompute the spherical to Cartesian transformation + # Pre-compute the spherical to Cartesian transformation # coefficients. num_ns = [] for l in range(max_angular + 1): diff --git a/anisoap/representations/ellipsoidal_density_projection.py b/anisoap/representations/ellipsoidal_density_projection.py index 6cd19ea..e4f2c72 100644 --- a/anisoap/representations/ellipsoidal_density_projection.py +++ b/anisoap/representations/ellipsoidal_density_projection.py @@ -1,4 +1,3 @@ -import sys import warnings from itertools import product @@ -10,11 +9,15 @@ ) from rascaline import NeighborList from scipy.spatial.transform import Rotation -from tqdm.auto import tqdm -from anisoap.representations.radial_basis import RadialBasis -from anisoap.utils.moment_generator import * -from anisoap.utils.spherical_to_cartesian import spherical_to_cartesian +from ..representations.radial_basis import RadialBasis +from ..utils.moment_generator import * +from ..utils.spherical_to_cartesian import spherical_to_cartesian + +try: + from tqdm.auto import tqdm +except ImportError: + tqdm = lambda x, **_: x def pairwise_ellip_expansion( @@ -95,12 +98,9 @@ def pairwise_ellip_expansion( leave=False, ) ): - frame_idx, i, j = ( - nl_sample["structure"], - nl_sample["first_atom"], - nl_sample["second_atom"], - ) - i_global = frame_to_global_atom_idx[frame_idx] + i + frame_idx, j = (nl_sample["structure"], nl_sample["second_atom"]) + + # i_global is not needed j_global = frame_to_global_atom_idx[frame_idx] + j r_ij = np.asarray( @@ -257,13 +257,11 @@ def contract_pairwise_feat(pair_ellip_feat, species, show_progress=False): block_samples = [] block_values = [] - for isample, sample in enumerate( - tqdm( - possible_block_samples, - disable=(not show_progress), - desc="Finding matching block samples", - leave=False, - ) + for sample in tqdm( + possible_block_samples, + disable=(not show_progress), + desc="Finding matching block samples", + leave=False, ): sample_idx = [ idx @@ -283,7 +281,7 @@ def contract_pairwise_feat(pair_ellip_feat, species, show_progress=False): # block_values has as many entries as samples satisfying (key, neighbor_species=ele). # When we iterate over neighbor species, not all (structure, center) would be present # Example: (0,0,1) might be present in a block with neighbor_species = 1 but no other pair block - # ever has (0,0,x) present as a sample- so (0,0) doesnt show up in a block_sample for all ele + # ever has (0,0,x) present as a sample- so (0,0) doesn't show up in a block_sample for all ele # so in general we have a ragged list of contract_blocks contract_blocks.append(block_values) @@ -304,11 +302,11 @@ def contract_pairwise_feat(pair_ellip_feat, species, show_progress=False): ) ) # Create storage for the final values - we need as many rows as all_block_samples, - # block.values.shape[1:] accounts for "components" and "properties" that are already part of the pair blocks - # and we dont alter these - # len(contract_blocks) - adds the additional dimension for the neighbor_species since we accumulated - # values for each of them as \sum_{j in ele} <|rho_ij> - # Thus - all_block_values.shape = (num_final_samples, components_pair, properties_pair, num_species) + # block.values.shape[1:] accounts for "components" and "properties" that + # are already part of the pair blocks and we dont alter these + # len(contract_blocks) - adds the additional dimension for the neighbor_species + # since we accumulated values for each of them as \sum_{j in ele} <|rho_ij> + # Thus - all_block_values.shape = (num_final_samples, components_pair, properties_pair, num_species) for iele, elem_cont_samples in enumerate( tqdm( @@ -319,15 +317,14 @@ def contract_pairwise_feat(pair_ellip_feat, species, show_progress=False): ) ): # This effectively loops over the species of the neighbors - # Now we just need to add the contributions to the final samples and values from this species to the right - # samples + # Now we just need to add the contributions to the final samples and + # values from this species to the right samples nzidx = [ i for i in range(len(all_block_samples)) if all_block_samples[i] in elem_cont_samples ] # identifies where the samples that this species contributes to, are present in the final samples - # print(apecies[ib],key, bb, all_block_samples) all_block_values[nzidx, :, :, iele] = contract_blocks[iele] new_block = TensorBlock( @@ -420,11 +417,12 @@ def __init__( raise ValueError( "radial_gaussian_width is set as an integer, which could cause overflow errors. Pass in float." ) - radial_hypers = {} - radial_hypers["radial_basis"] = radial_basis_name.lower() # lower case - radial_hypers["radial_gaussian_width"] = radial_gaussian_width - radial_hypers["max_angular"] = max_angular - self.radial_basis = RadialBasis(**radial_hypers) + radial_hypers = { + "radial_gaussian_width": radial_gaussian_width, + } + self.radial_basis = RadialBasis( + radial_basis_name.lower(), max_angular, **radial_hypers + ) self.num_ns = self.radial_basis.get_num_radial_functions() self.sph_to_cart = spherical_to_cartesian(self.max_angular, self.num_ns) @@ -455,7 +453,7 @@ def transform(self, frames, show_progress=False, normalize=True): show_progress : bool Show progress bar for frame analysis and feature generation normalize: bool - Whether to perform Lowdin Symmetric Orthonormalization or not. Orthonormalization generally + Whether to perform Löwdin Symmetric Orthonormalization or not. Orthonormalization generally leads to better performance. Default: True. Returns ------- @@ -483,8 +481,6 @@ def transform(self, frames, show_progress=False, normalize=True): # Define variables determining size of feature vector coming from frames self.num_atoms_per_frame = np.array([len(frame) for frame in frames]) - num_particle_types = len(species) - # Initialize arrays in which to store all features self.feature_gradients = 0 diff --git a/anisoap/representations/radial_basis.py b/anisoap/representations/radial_basis.py index 9c1e208..cbbcd3a 100644 --- a/anisoap/representations/radial_basis.py +++ b/anisoap/representations/radial_basis.py @@ -1,7 +1,6 @@ import warnings import numpy as np -import scipy.linalg from metatensor import TensorMap from scipy.special import gamma @@ -33,22 +32,22 @@ def inverse_matrix_sqrt(matrix: np.array): 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: + An un-normalized GTO of order n is \phi_n = r^n * e^{-r^2/(2*\sigma^2)} + The square norm of the un-normalized 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 + square norm: The square norm of the un-normalized 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. + Computes the normalization prefactor of an un-normalized 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: @@ -123,20 +122,24 @@ def __init__(self, radial_basis, max_angular, **hypers): def get_num_radial_functions(self): return self.num_radial_functions - # For each particle pair (i,j), we are provided with the three quantities - # that completely define the Gaussian distribution, namely - # the pair distance r_ij, the rotation matrix specifying the orientation - # of particle j's ellipsoid, as well the the three lengths of the - # principal axes. - # Combined with the choice of radial basis, these completely specify - # the mathematical problem, namely the integral that needs to be - # computed, which will be of the form - # integral gaussian(x,y,z) * polynomial(x,y,z) dx dy dz - # This function deals with the Gaussian part, which is specified - # by a precision matrix (inverse of covariance) and its center. - # The current function computes the covariance matrix and the center - # for the provided parameters as well as choice of radial basis. def compute_gaussian_parameters(self, r_ij, lengths, rotation_matrix): + """ + For each particle pair (i,j), we are provided with the three quantities + that completely define the Gaussian distribution, namely + the pair distance r_ij, the rotation matrix specifying the orientation + of particle j's ellipsoid, as well the the three lengths of the + principal axes. + + Combined with the choice of radial basis, these completely specify + the mathematical problem, namely the integral that needs to be + computed, which will be of the form + integral gaussian(x,y,z) * polynomial(x,y,z) dx dy dz + + This function deals with the Gaussian part, which is specified + by a precision matrix (inverse of covariance) and its center. + The current function computes the covariance matrix and the center + for the provided parameters as well as choice of radial basis. + """ # Initialization center = r_ij diag = np.diag(1 / lengths**2) @@ -155,9 +158,9 @@ 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 + The overlap matrix is the first step to generating an orthonormal basis set of functions (Löwdin 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 + block use a different set of GTOs. Hence, we pre-compute 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. @@ -179,15 +182,15 @@ def calc_gto_overlap_matrix(self): def orthonormalize_basis(self, features: TensorMap): """ - Apply an in-place orthonormalization on the features, using Lodwin Symmetric Orthonormalization. + Apply an in-place orthonormalization on the features, using Löwdin 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: + An instructive example of Löwdin 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. + copy of features must be made before the function if you wish to retain the un-normalized values. radial_basis: An instance of RadialBasis Returns: diff --git a/anisoap/utils/metatensor_utils.py b/anisoap/utils/metatensor_utils.py index 9763eba..dd34496 100644 --- a/anisoap/utils/metatensor_utils.py +++ b/anisoap/utils/metatensor_utils.py @@ -259,14 +259,12 @@ def cg_combine( properties_a = ( block_a.properties ) # pre-extract this block as accessing a c property has a non-zero cost - samples_a = block_a.samples # and x_b for index_b, block_b in x_b.items(): lam_b = index_b["angular_channel"] order_b = index_b["order_nu"] properties_b = block_b.properties - samples_b = block_b.samples if other_keys_match is None: OTHERS = tuple(index_a[name] for name in other_keys_a) + tuple( @@ -298,9 +296,9 @@ def cg_combine( prop_ids_a = [] prop_ids_b = [] - for n_a, f_a in enumerate(properties_a): + for f_a in properties_a: prop_ids_a.append(tuple(f_a) + (lam_a,)) - for n_b, f_b in enumerate(properties_b): + for f_b in properties_b: prop_ids_b.append(tuple(f_b) + (lam_b,)) prop_ids_a = np.asarray(prop_ids_a) prop_ids_b = np.asarray(prop_ids_b) diff --git a/anisoap/utils/moment_generator.py b/anisoap/utils/moment_generator.py index 48ea899..e29dba8 100644 --- a/anisoap/utils/moment_generator.py +++ b/anisoap/utils/moment_generator.py @@ -1,9 +1,4 @@ import numpy as np -from numpy.testing import assert_allclose -from scipy.special import ( - comb, - gamma, -) # Define function to compute all moments for a single @@ -40,7 +35,7 @@ def compute_moments_single_variable(A, a, maxdeg): # Define function to compute all moments for a diagonal dilation matrix. -# The implementation focuses on conceptual simplicity, while sacrifizing +# The implementation focuses on conceptual simplicity, while sacrificing # memory efficiency. # To be more specific, the array "moments" allows us to access the value # of the moment simply as moments[n0,n1,n2]. @@ -54,7 +49,7 @@ def compute_moments_diagonal_inefficient_implementation( - principal_components: np.ndarray of shape (3,) Array containing the three principal components - a: np.ndarray of shape (3,) - Vectorial center of the trivariate Gaussian + Vectorial center of the tri-variate Gaussian - maxdeg: int Maximum degree for which the moments need to be computed. @@ -63,7 +58,7 @@ def compute_moments_diagonal_inefficient_implementation( moments[n0,n1,n2] is the (n0,n1,n2)-th moment of the Gaussian defined as .. math:: - = \int(x^{n_0} * y^{n_1} * z^{n_2}) * \exp(-0.5*(r-a).T@cov@(r-a)) dxdydz + = \int(x^{n_0} * y^{n_1} * z^{n_2}) * \exp(-0.5*(r-a).T@cov@(r-a)) dx dy dz \sum_{i=1}^{\\infty} x_{i} Note that the term "moments" in probability theory are defined for normalized Gaussian distributions. @@ -77,7 +72,7 @@ def compute_moments_diagonal_inefficient_implementation( # The advantage, however, is the simplicity in later use. moments = np.zeros((maxdeg + 1, maxdeg + 1, maxdeg + 1)) - # Precompute the single variable moments in x- y- and z-directions: + # Pre-compute the single variable moments in x- y- and z-directions: moments_x = compute_moments_single_variable(principal_components[0], a[0], maxdeg) moments_y = compute_moments_single_variable(principal_components[1], a[1], maxdeg) moments_z = compute_moments_single_variable(principal_components[2], a[2], maxdeg) @@ -103,7 +98,7 @@ def compute_moments_diagonal_inefficient_implementation( # Define function to compute all moments for a general dilation matrix. -# The implementation focuses on conceptual simplicity, while sacrifizing +# The implementation focuses on conceptual simplicity, while sacrificing # memory efficiency. # To be more specific, the array "moments" allows us to access the value # of the moment simply as moments[n0,n1,n2]. @@ -118,7 +113,7 @@ def compute_moments_inefficient_implementation(A, a, maxdeg): the orientation of the three principal axes, while D is a diagonal matrix whose three diagonal elements are the lengths of the principal axes. - a: np.ndarray of shape (3,) - Vectorial center of the trivariate Gaussian. + Vectorial center of the tri-variate Gaussian. - maxdeg: int Maximum degree for which the moments need to be computed. @@ -126,7 +121,7 @@ def compute_moments_inefficient_implementation(A, a, maxdeg): - The list of moments defined as .. math:: - = \int(x^{n_0} * y^{n_1} * z^{n_2}) * \exp(-0.5*(r-a).T@cov@(r-a)) dxdydz + = \int(x^{n_0} * y^{n_1} * z^{n_2}) * \exp(-0.5*(r-a).T@cov@(r-a)) dx dy dz \sum_{i=1}^{\\infty} x_{i} Note that the term "moments" in probability theory are defined for normalized Gaussian distributions. diff --git a/anisoap/utils/monomial_iterator.py b/anisoap/utils/monomial_iterator.py index fd797dc..a6d9371 100644 --- a/anisoap/utils/monomial_iterator.py +++ b/anisoap/utils/monomial_iterator.py @@ -1,6 +1,3 @@ -import numpy as np - - class TrivariateMonomialIndices: """ Class for generating an iterator object over all trivariate @@ -23,7 +20,6 @@ class TrivariateMonomialIndices: myiter = iter(TrivariateMonomialIndices(deg=2)) for idx, n0, n1, n2 in myiter: ... # do something with exponents (n0, n1, n2) - """ def __init__(self, deg): @@ -53,8 +49,6 @@ def __next__(self): def find_idx(self, exponents): """ - - Parameters ---------- exponents : 3-tuple (n0, n1, n2) @@ -63,7 +57,6 @@ def find_idx(self, exponents): Returns ------- The index of the tuple in the lexicographical order - """ assert n0 + n1 + n2 == self.deg return self.exponent_list.index(exponents) diff --git a/anisoap/utils/spherical_to_cartesian.py b/anisoap/utils/spherical_to_cartesian.py index 3c1237c..6ccb9fd 100644 --- a/anisoap/utils/spherical_to_cartesian.py +++ b/anisoap/utils/spherical_to_cartesian.py @@ -1,12 +1,11 @@ import numpy as np -import scipy from scipy.special import ( comb, factorial, factorial2, ) -from anisoap.utils import monomial_iterator +from ..utils import monomial_iterator # Here we are implementing recurrence of the form R_{l}^m = prefact_minus1* z * T_{l-1} + prefact_minus2* r2 * T_{l-2} # where R_l^m is a solid harmonic, when expressed on a monomial basis - R_l^m = \sum_{n0+n1+n2=l} T_{l}[n0,n1,n2] x^n0 y^n1 z^n2 @@ -96,7 +95,7 @@ def spherical_to_cartesian(lmax, num_ns): for l in range(lmax + 1): deg = l myiter = iter(monomial_iterator.TrivariateMonomialIndices(deg)) - for idx, n0, n1, n2 in myiter: + for _, n0, n1, n2 in myiter: a = prefact_minus1(l - 1) # elements corresponding to m in (-l+2, ... l-2) b = prefact_minus2(l - 1) # elements corresponding to m in (-l+1, .... l-1) @@ -119,7 +118,7 @@ def spherical_to_cartesian(lmax, num_ns): for n in range(1, num_ns[l]): deg = l + 2 * n # degree of polynomial myiter = iter(monomial_iterator.TrivariateMonomialIndices(deg)) - for idx, n0, n1, n2 in myiter: + for _, n0, n1, n2 in myiter: # Use recurrence relation to update if n0 >= 2: T[l][:, n, n0, n1, n2] += T[l][:, n - 1, n0 - 2, n1, n2] diff --git a/tests/test_ellipsoidal_density_projection.py b/tests/test_ellipsoidal_density_projection.py index 6a8de21..80e07eb 100644 --- a/tests/test_ellipsoidal_density_projection.py +++ b/tests/test_ellipsoidal_density_projection.py @@ -1,7 +1,4 @@ -import builtins - -import ase -import metatensor +import ase import numpy as np import pytest diff --git a/tests/test_spherical_to_cartesian.py b/tests/test_spherical_to_cartesian.py index e0745c8..e21f523 100644 --- a/tests/test_spherical_to_cartesian.py +++ b/tests/test_spherical_to_cartesian.py @@ -1,8 +1,8 @@ from math import factorial +from numpy.testing import assert_allclose import numpy as np import pytest -from numpy.testing import assert_allclose from anisoap.utils import ( TrivariateMonomialIndices,