Pixelization API + Cosmology
Simplified JAX Pixelization API
The pixelization API no longer needs the preload arrays to perform modeling in JAX, simplifying the use of pixelized source reconstructions.
For example, the API to set up a rectangular mesh now only needs its shape:
"""
__Mesh Shape__
The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the source,
set below to 28 x 28.
The `mesh_shape` must be fixed before modeling and cannot be a free parameter of the model, because JAX uses the
mesh shape to define static shaped arrays which use the mesh to reconstruct the source. For a rectangular
mesh, the same number of pixels must be used in the y and x directions.
__Edge Zeroing__
By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by
the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct
bright surface brightnesses, often because they fit residuals from the lens light subtraction.
For a rectangular mesh, the source code computes edge pixels internally using the known
pixels at the edge of the mesh.
"""
mesh_pixels_yx = 28
mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)
Pull Request: https://github.com/Jammy2211/PyAutoLens/pulls?q=is%3Apr+is%3Aclosed
Example: https://github.com/Jammy2211/autolens_workspace/blob/release/notebooks/imaging/features/pixelization/modeling.ipynb
Cosmology JAX
Astropy Cosmology imports all removed, with all Cosmology calculations now using in built PyAutoLens functions which fully support JAX:
JAX cored and elliptical NFW
Both now fully supported:
PyAutoLabs/PyAutoGalaxy#279
PyAutoLabs/PyAutoGalaxy#280
Adaptive Matern Kernel
Matern kernel with regularization which adapts to source's brightness implemented: PyAutoLabs/PyAutoArray#214
Tests on rectangular mesh show this is the best performing regularization scheme for the rectangular mesh.
Mesh Refactor:
The above API changes are part of a larger refactoring of the pixliezed source module in autoarray, with the mesh and mapper modules now much cleaner in terms of code:
Pull Request: #388
Convolution Refactor
The PSF convolution h as been moved to an operator module and also cleaned up, including better internal padding of data when light outside the mask blurs in.
Pull Requests:
github.com/Jammy2211/PyAutoArray/pulls?q=is%3Apr+is%3Aclosed
PyAutoLabs/PyAutoArray#222
NaN Handling
Improvements to NaN handling removing certain failures during lens modeling: