JAX-GalSim
JAX port of GalSim, for parallelized, GPU accelerated, and differentiable galaxy image simulations.
![CodSpeed Badge](https://img.shields.io/endpoint?url=https://codspeed.io/badge.json)
Disclaimer: This project is still in an early development phase, please use the reference GalSim implementation for any scientific applications.
Objective and Design
The goal of this library is to reimplement GalSim functionalities in pure JAX to allow for automatic differentiation, GPU acceleration, and batched computations.
Guiding Principles
- Strive to be a drop-in replacement for GalSim, i.e. provide a close match to the GalSim API.
- Each function/feature will be tested against the reference GalSim implementation.
- This package will aim to be a subset of GalSim (i.e. only contains functions with a reference GalSim implementation).
- Implementations should be easy to read and understand.
- Code should be pip-installable on any machine, no compilation required.
- Any notable differences between the JAX and reference implementations will be clearly documented.
Notable Differences
- JAX arrays are immutable. Thus, in-place operations on images are not possible and a new image is
returned instead. Also, the RNG classes cannot fill arrays and instead return new arrays.
- JAX arrays do not support all the kinds of views that numpy arrays support. This means that some
operations that work in GalSim may not work in JAX-GalSim (e.g., real views of complex images).
- JAX-GalSim uses a different random number generator than GalSim. This leads to different results in terms of both the
generated random numbers and in terms of which RNGs have stable discarding.
Contributing
Everyone can contribute to this project, please refer to the CONTRIBUTING.md document for details.
In short, to interact with the project you can:
Issues marked with contributions welcome or good first issue are particularly good places to start. These are great ways to learn more
about the inner workings of GalSim and how to code in JAX.
Current GalSim API Coverage
JAX-GalSim has implemented 22.6% of the GalSim API. See the list below for the supported APIs.
- galsim.Add
- galsim.AffineTransform
- galsim.Angle
- galsim.AngleUnit
- galsim.BaseDeviate
- galsim.BaseNoise
- galsim.BaseWCS
- galsim.BinomialDeviate
- galsim.Bounds
- galsim.BoundsD
- galsim.BoundsI
- galsim.Box
- galsim.CCDNoise
- galsim.CelestialCoord
- galsim.Chi2Deviate
- galsim.Convolution
- galsim.Convolve
- galsim.Cubic
- galsim.Deconvolution
- galsim.Deconvolve
- galsim.Delta
- galsim.DeltaFunction
- galsim.DeviateNoise
- galsim.Exponential
- galsim.FitsHeader
- galsim.FitsWCS
- galsim.GSFitsWCS
- galsim.GSObject
- galsim.GSParams
- galsim.GalSimBoundsError
- galsim.GalSimConfigError
- galsim.GalSimConfigValueError
- galsim.GalSimDeprecationWarning
- galsim.GalSimError
- galsim.GalSimFFTSizeError
- galsim.GalSimHSMError
- galsim.GalSimImmutableError
- galsim.GalSimIncompatibleValuesError
- galsim.GalSimIndexError
- galsim.GalSimKeyError
- galsim.GalSimNotImplementedError
- galsim.GalSimRangeError
- galsim.GalSimSEDError
- galsim.GalSimUndefinedBoundsError
- galsim.GalSimValueError
- galsim.GalSimWarning
- galsim.GammaDeviate
- galsim.Gaussian
- galsim.GaussianDeviate
- galsim.GaussianNoise
- galsim.Image
- galsim.ImageCD
- galsim.ImageCF
- galsim.ImageD
- galsim.ImageF
- galsim.ImageI
- galsim.ImageS
- galsim.ImageUI
- galsim.ImageUS
- galsim.Interpolant
- galsim.InterpolatedImage
- galsim.JacobianWCS
- galsim.Lanczos
- galsim.Linear
- galsim.Moffat
- galsim.Nearest
- galsim.OffsetShearWCS
- galsim.OffsetWCS
- galsim.PhotonArray
- galsim.Pixel
- galsim.PixelScale
- galsim.PoissonDeviate
- galsim.PoissonNoise
- galsim.Position
- galsim.PositionD
- galsim.PositionI
- galsim.Quintic
- galsim.Sensor
- galsim.Shear
- galsim.ShearWCS
- galsim.SincInterpolant
- galsim.Spergel
- galsim.Sum
- galsim.TanWCS
- galsim.Transform
- galsim.Transformation
- galsim.UniformDeviate
- galsim.VariableGaussianNoise
- galsim.WeibullDeviate
- galsim.fits.closeHDUList
- galsim.fits.readCube
- galsim.fits.readFile
- galsim.fits.readMulti
- galsim.fits.write
- galsim.fits.writeFile
- galsim.fitswcs.CelestialWCS
- galsim.noise.addNoise
- galsim.noise.addNoiseSNR
- galsim.random.permute
- galsim.utilities.g1g2_to_e1e2
- galsim.utilities.horner
- galsim.utilities.printoptions
- galsim.utilities.unweighted_moments
- galsim.utilities.unweighted_shape
- galsim.wcs.EuclideanWCS
- galsim.wcs.LocalWCS
- galsim.wcs.UniformWCS
_Note: The coverage list is generated automatically by the scripts/update_api_coverage.py
script. To update it, run python scripts/update_api_coverage.py
from the root of the repository._