Conversation
Switches out `numpy` for `jax.numpy`.
When using JAX don't import Numpy
This adds the MGE code needed for the Sersic profile.
Add small radial offset to these two profiles to remove the `nan` values at the center. For the power law profile use the expansion of the `hyp2f1` function given by Tessore and Metcalf 2015 (eqn 29).
Update doc (merge from jax_wrapper branch)
This change makes the MGE profile work better for near-circular mass profiles. This is done by using more terms of the continued fraction for large values of `|z|^2`. This makes the auto grad more accurate and not evaluate to `0` when it should not.
Feature/jax remove preload
Main to JAX
Needed to make the `autolens.Tracer` example work.
Revert "Main to JAX"
Revert "Revert "Main to JAX""
Update functions for Tracer to be jax safe
Update Faddeeva function
These changes add the ability to use JAX to find the jacobian of deflection angles. To do this it first needs a version of the deflection angle function that takes in two scalars and returns a 2 element vector. From this `jax.jacfwd` is used to get the 2x2 jacobian for a single (y,x) value. This function is vectorized to take in a set of (y,x) values (of arbitrary shape) and return a (*shape, 2, 2) array of jacobian values. The final step is to pass the value from a `grid` object into this function.
Use auto diff to get at jacobian of deflection angles
This overhauls the current critical curve method from a grid search to an application of Newton's method for zero finding. The steps are as follows: 1) Create a set of `n_points` initial points in a circle of radius `init_r` and centred on `init_centre` 2) Apply `n_steps` of Newton's method to these points in the "radial" direction only (i.e. keeping angle fixed). Jax's auto differentiation is used to find the radial derivatives of the eigen value function for this step. 3) Filter the results and only keep point that have their eigen value `threshold` of 0. This step is needed because Newton's method will fail when the radial derivative is equal to zero (i.e. near a max or min of the eigen value). No underlying grid is needed for this method, but the quality of the results are dependent on the initial circle of points. For best results the circle should be drawn around *each* max/min of the eigen value function. This is typically the centre of each mass in the profile.
Calculate critical curves using JAX
Feature/jax merge
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
3 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.
Removes PyLops dependency and removes it from docs.