From aa0102aa4fcc4a444ffb4708c9ce61329d0934cb Mon Sep 17 00:00:00 2001 From: "Rose K. Cersonsky" <47536110+rosecers@users.noreply.github.com> Date: Mon, 13 Apr 2026 10:31:57 -0500 Subject: [PATCH 1/4] Making outputs from power spectrum standardized --- .../representations/ellipsoidal_density_projection.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/anisoap/representations/ellipsoidal_density_projection.py b/anisoap/representations/ellipsoidal_density_projection.py index 07d3645..62e21c8 100644 --- a/anisoap/representations/ellipsoidal_density_projection.py +++ b/anisoap/representations/ellipsoidal_density_projection.py @@ -726,6 +726,8 @@ def power_spectrum( "numbers", ] + n_particles = 0 + types_and_centers = [] # Checks if the sample contains all necessary information for computation of power spectrum for index, frame in enumerate(frames): array = frame.arrays @@ -736,6 +738,9 @@ def power_spectrum( ) if "quaternion" in array: raise ValueError(f"frame should contain c_q rather than quaternion") + for i in range(len(frame)): + n_particles += 1 + types_and_centers.append((index, i)) mvg_coeffs = self.transform( frames, show_progress=show_progress, rust_moments=rust_moments @@ -761,4 +766,8 @@ def power_spectrum( ) return x_asoap_raw else: - return mvg_nu2 + n_feat = mvg_nu2.block().values.squeeze().shape[1] + x_asoap_raw = np.zeros((n_particles, n_feat)) + for v, t in zip(mvg_nu2.block(0).values, mvg_nu2.block(0).samples): + x_asoap_raw[types_and_centers.index((t["type"], t["center"]))] = v + return x_asoap_raw From 60f5d094ab55212923c4a3a358c90d299604ee24 Mon Sep 17 00:00:00 2001 From: Arthur Lin Date: Mon, 13 Apr 2026 17:24:18 -0500 Subject: [PATCH 2/4] add logic for giving np.array outputs and require specification if tensormap is wanted --- .../ellipsoidal_density_projection.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/anisoap/representations/ellipsoidal_density_projection.py b/anisoap/representations/ellipsoidal_density_projection.py index 62e21c8..87da8af 100644 --- a/anisoap/representations/ellipsoidal_density_projection.py +++ b/anisoap/representations/ellipsoidal_density_projection.py @@ -683,7 +683,7 @@ def transform( return features def power_spectrum( - self, frames, mean_over_samples=True, show_progress=False, rust_moments=True + self, frames, mean_over_samples=True, show_progress=False, rust_moments=True, return_tensormap=False, ): """Helper function to compute the power spectrum of AniSOAP @@ -711,6 +711,7 @@ def power_spectrum( """ + assert not (mean_over_samples and return_tensormap) # if mean over samples == return tmap then error out - return tmap can only be true when mean-over-samples is false. # Initialize the Clebsch Gordan calculator for the angular component. mycg = ClebschGordanReal(self.max_angular) @@ -725,9 +726,6 @@ def power_spectrum( "positions", "numbers", ] - - n_particles = 0 - types_and_centers = [] # Checks if the sample contains all necessary information for computation of power spectrum for index, frame in enumerate(frames): array = frame.arrays @@ -738,9 +736,6 @@ def power_spectrum( ) if "quaternion" in array: raise ValueError(f"frame should contain c_q rather than quaternion") - for i in range(len(frame)): - n_particles += 1 - types_and_centers.append((index, i)) mvg_coeffs = self.transform( frames, show_progress=show_progress, rust_moments=rust_moments @@ -765,9 +760,7 @@ def power_spectrum( [block.values.squeeze() for block in x_asoap_raw.blocks()] ) return x_asoap_raw + elif return_tensormap: + return mvg_nu2 else: - n_feat = mvg_nu2.block().values.squeeze().shape[1] - x_asoap_raw = np.zeros((n_particles, n_feat)) - for v, t in zip(mvg_nu2.block(0).values, mvg_nu2.block(0).samples): - x_asoap_raw[types_and_centers.index((t["type"], t["center"]))] = v - return x_asoap_raw + return mvg_nu2.block().values.squeeze() From 68f84e809c9a87f6e22b059c3ccd873b1ce2d07b Mon Sep 17 00:00:00 2001 From: Arthur Lin Date: Mon, 13 Apr 2026 17:26:13 -0500 Subject: [PATCH 3/4] lint --- .../representations/ellipsoidal_density_projection.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/anisoap/representations/ellipsoidal_density_projection.py b/anisoap/representations/ellipsoidal_density_projection.py index 87da8af..e7044ec 100644 --- a/anisoap/representations/ellipsoidal_density_projection.py +++ b/anisoap/representations/ellipsoidal_density_projection.py @@ -683,7 +683,12 @@ def transform( return features def power_spectrum( - self, frames, mean_over_samples=True, show_progress=False, rust_moments=True, return_tensormap=False, + self, + frames, + mean_over_samples=True, + show_progress=False, + rust_moments=True, + return_tensormap=False, ): """Helper function to compute the power spectrum of AniSOAP @@ -711,7 +716,9 @@ def power_spectrum( """ - assert not (mean_over_samples and return_tensormap) # if mean over samples == return tmap then error out - return tmap can only be true when mean-over-samples is false. + assert not ( + mean_over_samples and return_tensormap + ) # if mean over samples == return tmap then error out - return tmap can only be true when mean-over-samples is false. # Initialize the Clebsch Gordan calculator for the angular component. mycg = ClebschGordanReal(self.max_angular) From d9895e5b271589b2ba0e68767db693300a76b74c Mon Sep 17 00:00:00 2001 From: Arthur Lin Date: Mon, 13 Apr 2026 17:38:42 -0500 Subject: [PATCH 4/4] add assert msg --- anisoap/representations/ellipsoidal_density_projection.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/anisoap/representations/ellipsoidal_density_projection.py b/anisoap/representations/ellipsoidal_density_projection.py index e7044ec..db1ef12 100644 --- a/anisoap/representations/ellipsoidal_density_projection.py +++ b/anisoap/representations/ellipsoidal_density_projection.py @@ -716,9 +716,10 @@ def power_spectrum( """ + # if mean over samples == return tmap then error out - return tmap can only be true when mean-over-samples is false. assert not ( mean_over_samples and return_tensormap - ) # if mean over samples == return tmap then error out - return tmap can only be true when mean-over-samples is false. + ), "cannot set return_tensormap to true if mean_over_samples is true" # Initialize the Clebsch Gordan calculator for the angular component. mycg = ClebschGordanReal(self.max_angular)