Merged
Conversation
feature/jax bool issue
These changes allow the PSF convolution to be calculated in a JAX friendly way that works with `jax.grad`. The method will take `image` and `blurring_image`, convert them from their 1D slim versions back into their 2D naive arrays (with masked points set to zero) and use `jax.scipy.signal.convolve` to do the convolution. A new keyword has been added to the function `jax_method` that can be either `direct` or `fft` to control how the convolution is done in JAX. Typically if the PSF kernel is more than about 5x5 the FFT method will be faster. Because of this `fft` is the default. Note: with the FFT some more speed might be gained by pre-caching the FFT of the PSF kernel as that is fixed and does not need to be re-computed every time, but this is an optimization that can be put in place at a later time once we know the current function works as expected.
Jax needs the base arrays passed in to work.
These files are all used by the Sersic profile's MGE deflection angles. This wholesale replacement of numpy for jax.numpy solves many of the issues around the array wrappers (e.g. "Grid2d"). The `numpy_wrapper` is just a simple file that can be imported to provide the correct `np` and an optional `use_jax` bool in case further `if` blocks are needed. Thankfully, with the full replacement much of the control flow `if` blocks seem to work without modification!
Adjust numpy to jax.numpy imports for files used for the MGE Sersic
Feature/remove convolver
rhayes777
approved these changes
Apr 9, 2025
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Doing whatever is possible to make unit tests pass, only modules remaining are interferometry and inversion which require dedicated care in their own PR.