Allow use of a subset of lens/source/psf models #3
Conversation
This file contains 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
…os. Need to do some performance testing
… nfw functions and the final translation to lenstronomy functions
…e dictionaries. Still needs to be tested and nfw functions need to be written.
…ittle more testing, but looking pretty good
…the some checkpoints are always kept
…s functions and no bugs found.
… the input pipeline will need some major refactor moving forward
…xible configuration file. Found some major bugs in the fixed configuration that were likely leading to problems with detecting sigma_sub. train.py still need a rework
…logy_params was being treated
…running smoothly after running tests on gpus. Signal should be very strong of the substructure now, so hopefully training should be smooth.
…off of desired truth outputs, and made sure that truncation radius can never return 0.
…calculations. Should speed up image simulations meaningfully.
… Substructure draws are similar, but images produced are different.
swagnercarena
force-pushed
the
model_selection
branch
from
dc638c5
to
e268647
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
This will allow us to implement new lens, source, or psf models without slowing down existing calculations.
Currently, to avoid parallelism-breaking conditionals (if lensmodel is NFW, then...), jaxstronomy evaluates the effect of every possible lens/source/psf model on the image, then selects the right output using
jax.lax.switchand the user-providedmodel_index.This PR adds optional arguments to the
image_simulationfunctions to specify the model lists in which themodel_indexlookups/switches occur. Lens light models obey thesource_modelsargument, since lens light models are source models (maybe we should call them light models or light profiles instead).Some new test code checks that you can change the order / amount of specified models for both lens, source, and PSF. The tests pass on my laptop GPU (assuming #2 is merged first; I can rebase if it helps you review).