Skip to content

Commit

Permalink
Added timer codes for benchmarking
Browse files Browse the repository at this point in the history
Implemented Rust FFI for moments and caching for CGR.

"Committing Rust FFI-related files"

Changed name of csv file to something shorter.

Changed the file system so that the file can be run from the root with python3 -m tests.ex_input

Moved all configs to ex_inputs.py file and ran diff check on results from two languages.

Implemented max-16-iter rule check for time mode

Fixed a bug that occurred when code_timer in ex_input.py was set to None.

Deleted unnecessary folder, equistore.

Changed the script to perform comparison between different versions

Updating isort on __init__

Update tests.yml to include coverage (#8)

Update tests.yml

Update tests.yml

Update tests.yml

Update README.md

Update tests.yml (#9)

* Update tests.yml

* Update tests.yml

Update tests.yml (#9)

* Update tests.yml

* Update tests.yml

Adding new tests for EDP (#7)

* Adding new tests for EDP

* Making requisite changes for equistore compatibility

* improved code coverage to 99% by testing for 3 additional cases: show_progress=True, multiple frames, and matrix rotations

* pass the linter

---------

Co-authored-by: Arthur Lin <[email protected]>

5 incorporate normalization factors (#6)

* Added (and made default behavior) the ability to orthonormalize features that use the GTO basis.
* This involves normalizing the features properly, creating an overlap matrix with orthogonal GTOs, and orthogonalizing the features.
* Added relevant tests to test new orthonormality functionality
* Added a jupyter notebook displaying how Lowdin Orthonormalization works (on a small gto basis set).
---------

Co-authored-by: Rose K. Cersonsky <[email protected]>
Co-authored-by: Arthur Lin <[email protected]>

minor changes to accomodate new equistore api

Updating to be in line with new equistore formatting (#15)

Adding progress bar for sanity sake (#14)

added warning for passing in integer, and cast any int arguments to f… (#13)

* Raise an error when a float is passed for radial_gaussian_width, and added a unit test to ensure error is raised

---------

Co-authored-by: Arthur

Added code that allows multiple iterations per parameter set. Now uses tqdm to keep track of progress.

Removed all internal timing code. The version selection feature still remains for testing purposes.

Implemented v2, in which loop for pairwise_ellip_expansion changed (further testing needed). Also, minor changes to output file.

Cross-platform support? Needs further testing

Added mac move command

Fixed file extension error on Mac

Moved caching logic to CGR class itself, instead of passing it as a keyword argument.

Removed unnecessary keyword argument from single_pass function

minor syntax changes for 3.9

added test ellipsoid trimers for performance testing

Added subprocess shell command to automatically run makefile.

Running makefile automatically while pip install version 2.

Added ell-trimers.xyz file.

Cross platform support... hopefully. Version 3.

Added MANIFEST.in file

Changed l_cut to lcut in cg_combine.

Fixed the issue of v0's output differing from the original implementation.

Some minor cleanup before some more experimental changes

Some code cleanup, mostly on the Rust side.

Updated equistore to metatensor. Requires changing _classes.py of metatensor.

Updated metatensor and rascaline to latest version. Changed import statements, and no _classes.py modification necessary.

Removed unused (for now) Rust dependencies

Some basic modification for test.

Added code to test running time depending on frame numbers.

Added code to test running time depending on frame numbers.

Changed implementation to determine repetition number with argv

Changed CGRCacheList replacement technique from FIFO to Clock (and added documentation)

Fixed a detail in Clock algorithm (initial insertion should have replacement flag = 0)

Fixed uninitialized _keys in cyclic_list.py

Fixed uninitialized _keys in cyclic_list.py

Updated metatensor field names to the most recent changes

Removing equistore in favor of metatensor (#20)

* Removing equistore in favor of metatensor

* fixing metatensor.core

Adding subtract self to the neighborlist generator (#19)

Switching list comp (#24)

Switching the code to do the list comp once instead of within each loop

Fixed mismatch between Python and Rust compute_moment

Cleaned up unused imports and variables.

Fixed minor spelling error

Deleted unncessary space in monomial_iterator.py

Ran isort to satisfy the linter.

Changed Rust implementation such that it takes inverse and determinant from Numpy instead of calculating internally.

Deleted version control switches and ran isort to satisfy linter.

Ran isort again.

Fixed a bug that occurs when both error AND inverse are 0. Copied main's contract_pairwise_feat code.

Adding ability to override number of radial bases

manual override n and added test cases

removed extraneous arg

Linter

Linter

Pointing to cutoff radius properly

I messed with the history so now I'm restoring it

pass linter

removed error check restricting independent specification of max_radial and radial_gaussian_width

added tests to ensure constant n functionality is correct

linter

num_radial -> max_radial in edp.py

removed changes pertaining to coupling sigma and n to keep PR clean

restored original compute_gaussian_parameters function header

incorporated necessary changes to make constant n work for both orthonormalization and moments generation

pass linter

removed deprecated comment

made definition of maxradial consistent with num_n and updated tests

Update lint.yml (#25)

add diff to linter so we can see how files are failing the linter.

Used Black formatter
  • Loading branch information
YCC-ProjBackups authored and rosecers committed Dec 5, 2023
1 parent e0d42ef commit 916b2e7
Show file tree
Hide file tree
Showing 23 changed files with 21,595 additions and 168 deletions.
17 changes: 17 additions & 0 deletions MANIFEST.in
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
global-exclude */__pycache__/*
global-exclude .pyc
global-exclude .DS_STORE
global-exclude *.so

recursive-include anisoap/reference *
recursive-include anisoap/representations *
recursive-include anisoap/utils *
recursive-include anisoap/lib *

include anisoap_rust_lib/Cargo.toml
recursive-include anisoap_rust_lib/src *

include Cargo.*
include pyproject.toml
include AUTHORS
include LICENSE
1 change: 1 addition & 0 deletions anisoap/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from anisoap import (
lib,
representations,
utils,
)
Expand Down
20 changes: 20 additions & 0 deletions anisoap/lib/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
import sys

import numpy

from .anisoap_rust_lib import compute_moments_ffi

#! Delete if not necessary -- need more testing on Windows OS.
if sys.platform.find("win") >= 0:
sys.path.append("c:\\Python\\DLLs")


def compute_moments(
dil_mat: numpy.ndarray, gau_cen: numpy.ndarray, max_deg: int
) -> numpy.ndarray:
"""
A simple wrapper around the Rust implementation for better integration with Python.
"""
return compute_moments_ffi(
numpy.linalg.inv(dil_mat), gau_cen, max_deg, numpy.linalg.det(dil_mat)
)
18 changes: 7 additions & 11 deletions anisoap/reference/projection_coefficients.py
Original file line number Diff line number Diff line change
@@ -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:
Expand Down Expand Up @@ -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):
Expand All @@ -65,7 +61,7 @@ def __init__(
# Initialize the radial basis class
if radial_basis_name not in ["monomial", "gto"]:
raise ValueError(
f"{self.radial_basis} is not an implemented basis"
f"{radial_basis_name} is not an implemented basis"
". Try 'monomial' or 'gto'"
)
if radial_gaussian_width != None and radial_basis_name != "gto":
Expand Down
190 changes: 91 additions & 99 deletions anisoap/representations/ellipsoidal_density_projection.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,15 +12,22 @@
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 anisoap.lib import compute_moments

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

from anisoap.lib import compute_moments

def pairwise_ellip_expansion(
lmax,
neighbor_list,
species,
frame_to_global_atom_idx,
rotation_matrices,
ellipsoid_lengths,
Expand Down Expand Up @@ -75,83 +82,77 @@ def pairwise_ellip_expansion(
keys = np.asarray(neighbor_list.keys, dtype=int)
keys = [tuple(i) + (l,) for i in keys for l in range(lmax + 1)]
num_ns = radial_basis.get_num_radial_functions()
maxdeg = np.max(np.arange(lmax + 1) + 2 * np.array(num_ns))
for center_species in species:
for neighbor_species in species:
if (center_species, neighbor_species) in neighbor_list.keys:
values_ldict = {l: [] for l in range(lmax + 1)}
nl_block = neighbor_list.block(
species_first_atom=center_species,
species_second_atom=neighbor_species,
)

for isample, nl_sample in enumerate(
tqdm(
nl_block.samples,
disable=(not show_progress),
desc="Iterating samples for ({}, {})".format(
center_species, neighbor_species
),
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
j_global = frame_to_global_atom_idx[frame_idx] + j

r_ij = np.asarray(
[
nl_block.values[isample, 0],
nl_block.values[isample, 1],
nl_block.values[isample, 2],
]
).reshape(
3,
)
for center_species, neighbor_species in neighbor_list.keys:
nl_block = neighbor_list.block(
species_first_atom=center_species,
species_second_atom=neighbor_species,
)
# moved from original position
values_ldict = {l: [] for l in range(lmax + 1)}
for isample, nl_sample in enumerate(
tqdm(
nl_block.samples,
disable=(not show_progress),
desc="Iterating samples for ({}, {})".format(
center_species, neighbor_species
),
leave=False,
)
):
frame_idx, j = (
nl_sample["structure"],
nl_sample["second_atom"],
)

rot = rotation_matrices[j_global]
lengths = ellipsoid_lengths[j_global]
precision, center = radial_basis.compute_gaussian_parameters(
r_ij, lengths, rot
)
r_ij = np.asarray(
[
nl_block.values[isample, 0],
nl_block.values[isample, 1],
nl_block.values[isample, 2],
]
).reshape(
3,
)

moments = compute_moments_inefficient_implementation(
precision, center, maxdeg=maxdeg
)
for l in range(lmax + 1):
deg = l + 2 * (num_ns[l] - 1)
moments_l = moments[: deg + 1, : deg + 1, : deg + 1]
values_ldict[l].append(
np.einsum("mnpqr, pqr->mn", sph_to_cart[l], moments_l)
)

for l in tqdm(
range(lmax + 1),
disable=(not show_progress),
desc="Accruing lmax",
leave=False,
):
block = TensorBlock(
values=np.asarray(values_ldict[l]),
samples=nl_block.samples, # as many rows as samples
components=[
Labels(
["spherical_component_m"],
np.asarray([list(range(-l, l + 1))], np.int32).reshape(
-1, 1
),
)
],
properties=Labels(
["n"],
np.asarray(list(range(num_ns[l])), np.int32).reshape(-1, 1),
),
# moved from original position
j_global = frame_to_global_atom_idx[frame_idx] + j
rot = rotation_matrices[j_global]
lengths = ellipsoid_lengths[j_global]

precision, center = radial_basis.compute_gaussian_parameters(
r_ij, lengths, rot
)
moments = compute_moments(precision, center, lmax + np.max(num_ns))

for l in range(lmax + 1):
deg = l + 2 * (num_ns[l] - 1)
moments_l = moments[: deg + 1, : deg + 1, : deg + 1]
values_ldict[l].append(
np.einsum("mnpqr, pqr->mn", sph_to_cart[l], moments_l)
)

for l in tqdm(
range(lmax + 1),
disable=(not show_progress),
desc="Accruing lmax",
leave=False,
):
block = TensorBlock(
values=np.asarray(values_ldict[l]),
samples=nl_block.samples, # as many rows as samples
components=[
Labels(
["spherical_component_m"],
np.asarray([list(range(-l, l + 1))], np.int32).reshape(-1, 1),
)
tensorblock_list.append(block)
],
properties=Labels(
["n"],
np.asarray(list(range(num_ns[l])), np.int32).reshape(-1, 1),
),
)
tensorblock_list.append(block)

pairwise_ellip_feat = TensorMap(
Labels(
Expand Down Expand Up @@ -230,7 +231,6 @@ def contract_pairwise_feat(pair_ellip_feat, species, show_progress=False):
]

if not len(sel_blocks):
# print(key, ele, "skipped") # this block is not found in the pairwise feat
continue
assert len(sel_blocks) == 1

Expand Down Expand Up @@ -266,13 +266,11 @@ def contract_pairwise_feat(pair_ellip_feat, species, show_progress=False):
l.append(idx)
indexed_sample_idx[tup] = l

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,
):
if sample in indexed_sample_idx:
sample_idx = indexed_sample_idx[tuple(sample)]
Expand All @@ -290,7 +288,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)
Expand All @@ -311,11 +309,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)

indexed_elem_cont_samples = {}
for i, val in enumerate(all_block_samples):
Expand Down Expand Up @@ -348,7 +346,6 @@ def contract_pairwise_feat(pair_ellip_feat, species, show_progress=False):
)

# 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(
Expand Down Expand Up @@ -431,7 +428,6 @@ def __init__(
# Currently, gradients are not supported
if compute_gradients:
raise NotImplementedError("Sorry! Gradients have not yet been implemented")
#

# Initialize the radial basis class
if radial_basis_name not in ["monomial", "gto"]:
Expand Down Expand Up @@ -485,7 +481,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
-------
Expand Down Expand Up @@ -513,8 +509,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

Expand Down Expand Up @@ -556,7 +550,6 @@ def transform(self, frames, show_progress=False, normalize=True):
pairwise_ellip_feat = pairwise_ellip_expansion(
self.max_angular,
self.nl,
species,
self.frame_to_global_atom_idx,
rotation_matrices,
ellipsoid_lengths,
Expand All @@ -567,7 +560,6 @@ def transform(self, frames, show_progress=False, normalize=True):

features = contract_pairwise_feat(pairwise_ellip_feat, species, show_progress)
if normalize:
normalized_features = self.radial_basis.orthonormalize_basis(features)
return normalized_features
return self.radial_basis.orthonormalize_basis(features)
else:
return features
Loading

0 comments on commit 916b2e7

Please sign in to comment.