grater_jax.disk_model.jax_model_wrappers
jax_model_wrappers.py
JAX-accelerated forward models for scattered-light disks.
This module provides JIT-compiled functions to generate disk model images using different combinations of disk models, dust distributions, scattering phase functions, and PSF models. Each function is specialized for particular choices of PSF/scattering implementations (e.g., Winnie PSF or spline-based SPF).
Functions
|
Get the generated disk model image given disk, scattering function, point spread function, stellar psf point spread function, and misceallaneous parameters along with the target image and error map. |
|
Get the generated disk model image given disk, scattering function, point spread function, stellar psf point spread function, and misceallaneous parameters along with the target image and error map. |
|
Get the generated disk model image given disk, scattering function, point spread function, stellar psf point spread function, and misceallaneous parameters along with the target image and error map. |
|
Get the generated disk model image given disk, scattering function, point spread function, stellar psf point spread function, and misceallaneous parameters along with the target image and error map. |
Classes
|
partial(func, *args, **keywords) - new function with partial application of the given arguments and keywords. |
- grater_jax.disk_model.jax_model_wrappers.jax_model(DiskModel, DistrModel, FuncModel, PSFModel, StellarPSFModel, disk_params, spf_params, psf_params, stellar_psf_params, throughput, distance=0.0, pxInArcsec=0.0, nx=140, ny=140, halfNbSlices=25, flux_scaling=1000000.0)
Get the generated disk model image given disk, scattering function, point spread function, stellar psf point spread function, and misceallaneous parameters along with the target image and error map. This function only applies when PSFModel != Winnie_PSF and FuncModel != InterpolatedUnivariateSpline_SPF.
- DiskModelclass (ScatteredLighDisk is the only supported disk model)
The disk model type
- DistrModelclass (DustEllipticalDistribution2PowerLaws is the only supported dust distribution model)
The dust distribution model type
- FuncModelclass (Can be found in disk_model/SLD_utils.py)
The scattering phase function model type
- PSFModelclass (Can be found in disk_model/SLD_utils.py)
The point spread function model type
- disk_paramsjnp.array
The corresponding parameter dictionary for the disk model, dictionary is made of (parameter name, parameter value) pairs.
- spf_paramsjnp.array
The corresponding parameter dictionary for the scattering phase function, dictionary is made of (parameter name, parameter value) pairs.
- psf_paramsjnp.array
The corresponding parameter dictionary for the point spread function, dictionary is made of (parameter name, parameter value) pairs.
- StellarPSFModelclass, optional
The scattering phase function model type, set to None be default indicating no stellar psf model.
- stellar_psf_paramsjnp.array, optional
The corresponding parameter dictionary for the on axis stellar psf model, dictionary is made of (parameter name, parameter value) pairs.
- target_imagenp.ndarray
The target image that the log likelihood is being computed for.
- err_mapnp.ndarray
The error map for the target image.
- throughputnp.ndarray
The throughput image to apply to the generated disk image.
- distancefloat
Distance to the star in pc (default 70.)
- pxInArcsecfloat
Pixel field of view in arcsec/px (default the SPHERE pixel scale 0.01225 arcsec/px)
- nxint
number of pixels along the x axis of the image (default 200)
- nyint
number of pixels along the y axis of the image (default 200)
- halfNbSlicesinteger
half number of distances along the line of sight
- flux_scalingfloat
Scaling factor for disk model.
- Returns:
Generated disk model image
- Return type:
jnp.ndarray
- grater_jax.disk_model.jax_model_wrappers.jax_model_winnie(DiskModel, DistrModel, FuncModel, winnie_psf, StellarPSFModel, disk_params, spf_params, stellar_psf_params, throughput, distance=0.0, pxInArcsec=0.0, nx=140, ny=140, halfNbSlices=25, flux_scaling=1000000.0)
Get the generated disk model image given disk, scattering function, point spread function, stellar psf point spread function, and misceallaneous parameters along with the target image and error map. This function only applies when PSFModel == Winnie_PSF and FuncModel != InterpolatedUnivariateSpline_SPF.
- DiskModelclass (ScatteredLighDisk is the only supported disk model)
The disk model type
- DistrModelclass (DustEllipticalDistribution2PowerLaws is the only supported dust distribution model)
The dust distribution model type
- FuncModelclass (Can be found in disk_model/SLD_utils.py)
The scattering phase function model type
- winnie_psfclass (Can be found in disk_model/winnie_class.py)
JWST off axis PSF, modeled after Winnie framework
- disk_paramsjnp.array
The corresponding parameter dictionary for the disk model, dictionary is made of (parameter name, parameter value) pairs.
- spf_paramsjnp.array
The corresponding parameter dictionary for the scattering phase function, dictionary is made of (parameter name, parameter value) pairs.
- StellarPSFModelclass, optional
The scattering phase function model type, set to None be default indicating no stellar psf model.
- stellar_psf_paramsjnp.array, optional
The corresponding parameter dictionary for the on axis stellar psf model, dictionary is made of (parameter name, parameter value) pairs.
- target_imagenp.ndarray
The target image that the log likelihood is being computed for.
- err_mapnp.ndarray
The error map for the target image.
- throughputnp.ndarray
The throughput image to apply to the generated disk image.
- distancefloat
Distance to the star in pc (default 70.)
- pxInArcsecfloat
Pixel field of view in arcsec/px (default the SPHERE pixel scale 0.01225 arcsec/px)
- nxint
number of pixels along the x axis of the image (default 200)
- nyint
number of pixels along the y axis of the image (default 200)
- halfNbSlicesinteger
half number of distances along the line of sight
- flux_scalingfloat
Scaling factor for disk model.
- Returns:
Generated disk model image
- Return type:
jnp.ndarray
- grater_jax.disk_model.jax_model_wrappers.jax_model_spline(DiskModel, DistrModel, FuncModel, PSFModel, StellarPSFModel, disk_params, spf_params, psf_params, stellar_psf_params, throughput, distance=0.0, pxInArcsec=0.0, nx=140, ny=140, halfNbSlices=25, flux_scaling=1000000.0, knots=array([1., 0.66666667, 0.33333333, 0., -0.33333333, -0.66666667, -1.]))
Get the generated disk model image given disk, scattering function, point spread function, stellar psf point spread function, and misceallaneous parameters along with the target image and error map. This function only applies when PSFModel != Winnie_PSF and FuncModel == InterpolatedUnivariateSpline_SPF.
- DiskModelclass (ScatteredLighDisk is the only supported disk model)
The disk model type
- DistrModelclass (DustEllipticalDistribution2PowerLaws is the only supported dust distribution model)
The dust distribution model type
- FuncModelclass (Can be found in disk_model/SLD_utils.py)
The scattering phase function model type
- PSFModelclass (Can be found in disk_model/SLD_utils.py)
The point spread function model type
- disk_paramsjnp.array
The corresponding parameter dictionary for the disk model, dictionary is made of (parameter name, parameter value) pairs.
- spf_paramsjnp.array
The corresponding parameter dictionary for the scattering phase function, dictionary is made of (parameter name, parameter value) pairs.
- psf_paramsjnp.array
The corresponding parameter dictionary for the point spread function, dictionary is made of (parameter name, parameter value) pairs.
- StellarPSFModelclass, optional
The scattering phase function model type, set to None be default indicating no stellar psf model.
- stellar_psf_paramsjnp.array, optional
The corresponding parameter dictionary for the on axis stellar psf model, dictionary is made of (parameter name, parameter value) pairs.
- target_imagenp.ndarray
The target image that the log likelihood is being computed for.
- err_mapnp.ndarray
The error map for the target image.
- throughputnp.ndarray
The throughput image to apply to the generated disk image.
- distancefloat
Distance to the star in pc (default 70.)
- pxInArcsecfloat
Pixel field of view in arcsec/px (default the SPHERE pixel scale 0.01225 arcsec/px)
- nxint
number of pixels along the x axis of the image (default 200)
- nyint
number of pixels along the y axis of the image (default 200)
- halfNbSlicesinteger
half number of distances along the line of sight
- flux_scalingfloat
Scaling factor for disk model.
- Returns:
Generated disk model image
- Return type:
jnp.ndarray
- grater_jax.disk_model.jax_model_wrappers.jax_model_spline_winnie(DiskModel, DistrModel, FuncModel, winnie_psf, StellarPSFModel, disk_params, spf_params, stellar_psf_params, throughput, distance=0.0, pxInArcsec=0.0, nx=140, ny=140, halfNbSlices=25, flux_scaling=1000000.0, knots=array([1., 0.66666667, 0.33333333, 0., -0.33333333, -0.66666667, -1.]))
Get the generated disk model image given disk, scattering function, point spread function, stellar psf point spread function, and misceallaneous parameters along with the target image and error map. This function only applies when PSFModel == Winnie_PSF and FuncModel == InterpolatedUnivariateSpline_SPF.
- DiskModelclass (ScatteredLighDisk is the only supported disk model)
The disk model type
- DistrModelclass (DustEllipticalDistribution2PowerLaws is the only supported dust distribution model)
The dust distribution model type
- FuncModelclass (Can be found in disk_model/SLD_utils.py)
The scattering phase function model type
- winnie_psfclass (Can be found in disk_model/winnie_class.py)
JWST off axis PSF, modeled after Winnie framework
- disk_paramsjnp.array
The corresponding parameter dictionary for the disk model, dictionary is made of (parameter name, parameter value) pairs.
- spf_paramsjnp.array
The corresponding parameter dictionary for the scattering phase function, dictionary is made of (parameter name, parameter value) pairs.
- psf_paramsjnp.array
The corresponding parameter dictionary for the point spread function, dictionary is made of (parameter name, parameter value) pairs.
- StellarPSFModelclass, optional
The scattering phase function model type, set to None be default indicating no stellar psf model.
- stellar_psf_paramsjnp.array, optional
The corresponding parameter dictionary for the on axis stellar psf model, dictionary is made of (parameter name, parameter value) pairs.
- target_imagenp.ndarray
The target image that the log likelihood is being computed for.
- err_mapnp.ndarray
The error map for the target image.
- throughputnp.ndarray
The throughput image to apply to the generated disk image.
- distancefloat
Distance to the star in pc (default 70.)
- pxInArcsecfloat
Pixel field of view in arcsec/px (default the SPHERE pixel scale 0.01225 arcsec/px)
- nxint
number of pixels along the x axis of the image (default 200)
- nyint
number of pixels along the y axis of the image (default 200)
- halfNbSlicesinteger
half number of distances along the line of sight
- flux_scalingfloat
Scaling factor for disk model.
- Returns:
Generated disk model image
- Return type:
jnp.ndarray