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

Implement SigmaG filtering in JAX #750

Merged
merged 2 commits into from
Dec 16, 2024
Merged

Implement SigmaG filtering in JAX #750

merged 2 commits into from
Dec 16, 2024

Conversation

jeremykubica
Copy link
Contributor

Change the numpy vectorized implementation of sigma G filtering to use JAX and vmap. This produced roughly a 2000x speedup on sigma-G filtering on timing tests on Baldur.

Also makes one behavior change. Previously sigma-G filtering with clipped_negative=True would filter out zero values of LH (but NOT all negative values). This does not make sense with the goal of sigma-G, so I have removed it.

Copy link
Collaborator

@wilsonbb wilsonbb left a comment

Choose a reason for hiding this comment

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

LGTM!

@jeremykubica jeremykubica merged commit 0e5198c into main Dec 16, 2024
2 checks passed
@jeremykubica jeremykubica deleted the sigma_g_jax branch December 16, 2024 13:08
@jeremykubica jeremykubica restored the sigma_g_jax branch December 16, 2024 20:56
@jeremykubica jeremykubica deleted the sigma_g_jax branch December 16, 2024 20:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants