Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Vectorize signal generator functions #252

Merged
merged 13 commits into from
Dec 18, 2024

Conversation

matt-graham
Copy link
Collaborator

@matt-graham matt-graham commented Dec 4, 2024

When running some benchmarks with large L I noticed that the signal generation functions in s2fft.utils.signal_generator become very slow for large L due to the nested for loops.

This switches to using a vectorized implementation. Currently implemented for generate_flm only.

EDIT: Now implemented for both generate_flm and generate_flmn.

Copy link

codecov bot commented Dec 4, 2024

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 96.14%. Comparing base (7cf80bb) to head (17055fb).
Report is 55 commits behind head on main.

Additional details and impacted files
@@            Coverage Diff             @@
##             main     #252      +/-   ##
==========================================
+ Coverage   96.07%   96.14%   +0.07%     
==========================================
  Files          31       31              
  Lines        3565     3581      +16     
==========================================
+ Hits         3425     3443      +18     
+ Misses        140      138       -2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@matt-graham matt-graham marked this pull request as ready for review December 11, 2024 13:28
@matt-graham
Copy link
Collaborator Author

matt-graham commented Dec 11, 2024

  • Update docstring with pseudo-code with equivalent loop-based implementation to give intuition for vectorized term
  • Correct m = 0 case when reality = False (should still be real)

Comment on lines 82 to +90
if precomps is None:
precomps = generate_precomputes(L, -mm, sampling, nside, L_lower)
precomps = generate_precomputes(
L=L,
spin=-mm,
sampling=sampling,
nside=nside,
forward=False,
L_lower=L_lower,
)
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is same change as fix in #256 to resolve #255. Applied here as otherwise test in tests/test_signal_generator.py::test_generate_flm was failing with L_lower != 0.

Comment on lines +165 to +174
elif n > 0:
# Generate independent complex coefficients for positive n
flmn[N - 1 + n, min_el:L, L - 1] = complex_normal(
rng, L - min_el, var=2
)
# For m = 0, n > 0
# flmn[N - 1 - n, el, L - 1] = (-1)**n * flmn[N - 1 + n, el, L - 1].conj
flmn[N - 1 - n, min_el:L, L - 1] = (-1) ** n * (
flmn[N - 1 + n, min_el:L, L - 1].conj()
)
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Compared to previous implementation, here for m = 0 and abs(n) > 0 we generate complex coefficients rather than real, as the conjugate symmetry condition does not seem to enforce reality in this case - not sure if this is correct though?

@matt-graham matt-graham added the enhancement New feature or request label Dec 13, 2024
Copy link
Collaborator

@CosmoMatt CosmoMatt left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice! We originally just implemented these functions for testing purposes so didn't bother optimising them, but as you say definitely worth quickly updating them now.

@matt-graham matt-graham merged commit f08c6bc into main Dec 18, 2024
13 checks passed
@matt-graham matt-graham deleted the mmg/vectorized-signal-generators branch December 18, 2024 13:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants