jaxtronomy.Util package

Submodules

jaxtronomy.Util.class_creator module

create_class_instances(lens_model_list=None, z_lens=None, z_source=None, z_source_convention=None, lens_redshift_list=None, lens_profile_kwargs_list=None, multi_plane=False, distance_ratio_sampling=False, cosmology_sampling=False, cosmology_model='FlatLambdaCDM', observed_convention_index=None, source_light_model_list=None, source_light_profile_kwargs_list=None, lens_light_model_list=None, lens_light_profile_kwargs_list=None, point_source_model_list=None, point_source_redshift_list=None, fixed_magnification_list=None, point_source_frame_list=None, additional_images_list=None, kwargs_lens_eqn_solver=None, source_deflection_scaling_list=None, source_redshift_list=None, cosmo=None, index_lens_model_list=None, index_source_light_model_list=None, index_lens_light_model_list=None, index_point_source_model_list=None, optical_depth_model_list=None, optical_depth_profile_kwargs_list=None, index_optical_depth_model_list=None, band_index=0, tau0_index_list=None, all_models=False, point_source_magnification_limit=None, decouple_multi_plane=False, kwargs_multiplane_model=None, kwargs_multiplane_model_point_source=None, tracer_source_model_list=None, tracer_source_band=0, tracer_partition=None, tracer_type='LINEAR')[source]
Parameters:
  • lens_model_list – list of strings indicating the type of lens models

  • z_lens – redshift of the deflector (for single lens plane mode, but only relevant when computing physical quantities)

  • z_source – redshift of source (for single source plane mode, or for multiple source planes the redshift of the point source). In regard to this redshift the reduced deflection angles are defined in the lens model.

  • z_source_convention – float, redshift of a source to define the reduced deflection angles of the lens models. If None, ‘z_source’ is used.

  • lens_redshift_list – None or list of floats in the same order of the lens_model_list

  • lens_profile_kwargs_list – list of dicts, keyword arguments used to initialize deflector profile classes in the same order of the lens_model_list. If any of the profile_kwargs are None, then that profile will be initialized using default settings.

  • multi_plane – bool, if True, computes the lensing quantities in multi-plane mode

  • distance_ratio_sampling – bool, if True, samples the distance ratios in multi-lens-plane

  • cosmology_sampling – bool, if True, samples the cosmology in multi-lens-plane

  • cosmology_model – string, name of the cosmology model to be used in the multi-lens-plane mode

  • observed_convention_index

  • source_light_model_list – list of strings indicating the type of source light models

  • source_light_profile_kwargs_list – list of dicts, keyword arguments used to initialize source light profile classes in the same order of the source_light_model_list. If any of the profile_kwargs are None, then that profile will be initialized using default settings.

  • lens_light_model_list – list of strings indicating the type of lens light models

  • lens_light_profile_kwargs_list – list of dicts, keyword arguments used to initialize lens light profile classes in the same order of the lens_light_model_list. If any of the profile_kwargs are None, then that profile will be initialized using default settings.

  • point_source_model_list – list of strings indicating the type of point source models

  • fixed_magnification_list – list of bool. Indicates which point source classes in the same order of point_source_model_list should have fixed magnification. Only relevant for the LENSED_POSITION point source type. If set to True, then “source_amp” is a parameter instead of “point_amp”, and the magnification is calculated from the lens models.

  • point_source_frame_list – list of list of ints, assigns each model in point_source_type_list a frame list. Only relevent for LENSED_POSITION. e.g. if point_source_type_list = [“UNLENSED”, “LENSED_POSITION”, “LENSED_POSITION”] with point_source_frame_list = [None, [0, 1, 2], [1, 2, 0, 1]], then the first LENSED_POSITION will have a frame list of [0, 1, 2] and the second LENSED_POSITION will have a frame list of [1, 2, 0, 1]. See docstring of point_source_frame_list in PSBase for further details.

  • additional_images_list – list of bool. Indicates which point source classes in the same order of the point_source_model_list should use the lens equation solver to solve for additional images. Only relevant for the LENSED_POSITION point source type.

  • kwargs_lens_eqn_solver – keyword arguments specifying the numerical settings for the lens equation solver see LensEquationSolver() class for details

  • source_deflection_scaling_list – List of floats for each source ligth model (optional, and only applicable for single-plane lensing. The factors re-scale the reduced deflection angles described from the lens model. =1 means identical source position as without this option. This option enables multiple source planes. The geometric difference between the different source planes needs to be pre-computed and is cosmology dependent.

  • source_redshift_list – list of redshifts for the source model profiles in the same order of the source_light_model_list

  • cosmo – astropy.cosmology instance

  • index_lens_model_list – list of list of ints, indicating which lens models are in each band. For example [[0, 2], [1, 3]] indicates that band 0 uses lens models 0 and 2, and band 1 uses lens models 1 and 3 from the lens_model_list

  • index_source_light_model_list – list of list of ints, indicating which source light models are in each band.

  • index_lens_light_model_list – optional, list of list of all model indexes for each modeled band

  • index_point_source_model_list – optional, list of list of all model indexes for each modeled band

  • optical_depth_model_list – list of strings indicating the optical depth model to compute (differential) extinctions from the source

  • optical_depth_profile_kwargs_list – list of dicts, keyword arguments used to initialize light model profile classes in the same order of the optical_depth_model_list. If any of the profile_kwargs are None, then that profile will be initialized using default settings.

  • index_optical_depth_model_list – list of list of ints, indicates which optical depth models are in each band.

  • band_index – int, index of band to consider. Has an effect if only partial models are considered for a specific band

  • tau0_index_list – list of integers of the specific extinction scaling parameter tau0 for each band

  • all_models – bool, if True, will make class instances of all models ignoring potential keywords that are excluding specific models as indicated.

  • point_source_magnification_limit – float >0 or None, if set and additional images are computed, then it will cut the point sources computed to the limiting (absolute) magnification

  • decouple_multi_plane – bool; if True, creates an instance of MultiPlaneDecoupled

  • kwargs_multiplane_model – keyword arguments used to create an instance of MultiPlaneDecoupled if decouple_multi_plane is True

  • kwargs_multiplane_model_point_source – keyword arguments used to create an option MultiPlaneDecoupled class for the lensed point source to be treated separately from the rest of the imaging data

  • tracer_source_model_list – list of tracer source models (not used in this function)

  • tracer_source_band – integer, list index of source surface brightness band to apply tracer model to

  • tracer_partition (None or list) – in case of tracer models for specific sub-parts of the surface brightness model [[list of light profiles, list of tracer profiles], [list of light profiles, list of tracer profiles], […], …]

  • tracer_type – string with options ‘LINEAR’ or ‘LOG’, to determine how tracers are summed between components

Returns:

lens_model_class, source_model_class, lens_light_model_class, point_source_class, extinction_class

create_image_model(kwargs_data, kwargs_psf, kwargs_numerics, kwargs_model, image_likelihood_mask=None)[source]
Parameters:
  • kwargs_data – ImageData keyword arguments

  • kwargs_psf – PSF keyword arguments

  • kwargs_numerics – numerics keyword arguments for Numerics() class

  • kwargs_model – model keyword arguments

  • image_likelihood_mask – image likelihood mask (same size as image_data with 1 indicating being evaluated and 0 being left out)

Returns:

ImageModel() instance

create_im_sim(multi_band_list, multi_band_type, kwargs_model, bands_compute=None, image_likelihood_mask_list=None, band_index=0, kwargs_pixelbased=None, linear_solver=True)[source]
Parameters:
  • multi_band_list – list of [[kwargs_data, kwargs_psf, kwargs_numerics], [], ..]

  • multi_band_type – string, option when having multiple imaging data sets modelled simultaneously. Options are: - ‘single-band’: single band - ‘multi-linear’ and joint-linear are not supported in jaxtronomy

  • kwargs_model – model keyword arguments

  • bands_compute – (optional), bool list to indicate which band to be included in the modeling

  • image_likelihood_mask_list – list of image likelihood mask (same size as image_data with 1 indicating being evaluated and 0 being left out)

  • band_index – integer, index of the imaging band to model (only applied when using ‘single-band’ as option)

  • kwargs_pixelbased – keyword arguments with various settings related to the pixel-based solver (see SLITronomy documentation)

  • linear_solver – bool, if True (default) fixes the linear amplitude parameters ‘amp’ (avoid sampling) such that they get overwritten by the linear solver solution.

Returns:

MultiBand class instance

jaxtronomy.Util.herm_util module

eval_hermite(n, x)[source]

Equivalent to scipy.special.eval_hermite(n, x).

Parameters:
  • n – int, order of hermite polynomial to be evaluated

  • x – array-like, coordinates to evaluate hermite polynomial

Returns:

array with same shape as x containing H_n(x)

hermval(x, c)[source]

Equivalent to numpy.polynomial.hermite.hermval when the input c is 1 dimensional.

Parameters:
  • x – array-like, coordinates to evaluate hermite polynomials

  • c – array-like, with dimension equal to 1

Returns:

array with the same shape as x containing hermval(x)

jaxtronomy.Util.hyp2f1_util module

hyp2f1_series(a, b, c, z, nmax=50)[source]

This computation is based off of the standard series expansion of hyp2f1. The recurrence relation between successive terms in the sum is well known:

A_0 = 1 A_i = z * [(i - 1 + a)(i - 1 + b) / i(i - 1 + c)] A_{i-1}

The conditions required for this series to converge are: 1) |z| < 1 2) c is not a non-positive integer (i.e. c != 0, -1, -2, …) 3) If Re(c - a - b) > 0, then the series converges on |z| = 1 as well (but very slowly)

hyp2f1_continuation(a, b, c, z, nmax=50)[source]

This implementation is based off of Bühring’s analytic continuation formula Equation 4, from J.L. Lopez and N.M. Temme, “New series expansions of the Gauss hypergeometric function”, Adv Comput Math 39, 349-365 (2013). https://link.springer.com/content/pdf/10.1007/s10444-012-9283-y.pdf

The conditions required for this series to converge are: 1) |z-1/2| > 1/2. 2) b - a is not an integer

hyp2f1_lopez_temme_8(a, b, c, z, nmax=75)[source]

Equation 8 from J.L. Lopez and N.M. Temme, “New series expansions of the Gauss hypergeometric function”, Adv Comput Math 39, 349-365 (2013). https://link.springer.com/content/pdf/10.1007/s10444-012-9283-y.pdf

This series expansion converges whenever Re(z) < 1, and does not have any restrictions on a, b, or c. The downside is that this series converges slowly for |z| -> infty, so nmax needs to be higher for an accurate result. This series expansion coverges a bit faster than the standard series expansion when z is in the unit disk.

hyp2f1_near_one(a, b, c, z, nmax=50)[source]

This implementation is based off of equation 15.3.6 in Abramowitz and Stegun. This transformation formula allows for a calculation of hyp2f1 for points near.

z = 1 (where other iterative computation schemes converge slowly) by transforming z to 1 - z.

The conditions required for this series to converge are: 1) |1-z| < 1 2) z is not purely real such that z >= 1 3) c - a - b is not an integer

hyp2f1(a, b, c, z)[source]

This function looks at where z is located on the complex plane and chooses the appropriate hyp2f1 function to use in the interest of maintaining optimal runtime and accuracy.

If the user already knows which hyp2f1 function to use, this step can be skipped and the user can directly call the appropriate hyp2f1 function. If the input z is an array with values in different regions on the complex plane, this function MUST be used so that the appropriate hyp2f1 function is used for each value.

jaxtronomy.Util.image_util module

add_layer2image(grid2d, x_pos, y_pos, kernel, order=1)[source]

Adds a kernel on the grid2d image at position x_pos, y_pos with an interpolated subgrid pixel shift of order=order.

Parameters:
  • grid2d – 2d pixel grid (i.e. image)

  • x_pos – x-position center (pixel coordinate) of the layer to be added

  • y_pos – y-position center (pixel coordinate) of the layer to be added

  • kernel – the layer to be added to the image

  • order – interpolation order for sub-pixel shift of the kernel to be added

Returns:

image with added layer, cut to original size

add_layer2image_int(grid2d, x_pos, y_pos, kernel)[source]

Adds a kernel on the grid2d image at position x_pos, y_pos at integer positions of pixel.

Parameters:
  • grid2d – 2d pixel grid (i.e. image)

  • x_pos – x-position center (pixel coordinate) of the layer to be added

  • y_pos – y-position center (pixel coordinate) of the layer to be added

  • kernel – the layer to be added to the image

Returns:

image with added layer

re_size(image, factor=1)[source]

Re-sizes image with nx x ny to nx/factor x ny/factor.

Parameters:
  • image – 2d image with shape (nx,ny)

  • factor – integer >=1

Returns:

jaxtronomy.Util.kernel_util module

Routines that manipulate convolution kernels.

estimate_amp(data, x_pos, y_pos, psf_kernel)[source]

Estimates the amplitude of a point source located at x_pos, y_pos.

Parameters:
  • data

  • x_pos

  • y_pos

  • psf_kernel

Returns:

jaxtronomy.Util.param_util module

cart2polar(x, y, center_x=0, center_y=0)[source]

Transforms cartesian coords [x,y] into polar coords [r,phi] in the frame of the lens center.

Parameters:
  • x (array of size (n)) – set of x-coordinates

  • y (array of size (n)) – set of x-coordinates

  • center_x (float) – rotation point

  • center_y (float) – rotation point

Returns:

array of same size with coords [r,phi]

polar2cart(r, phi, center)[source]

Transforms polar coords [r,phi] into cartesian coords [x,y] in the frame of the lense center.

Parameters:
  • r (array of size n or float) – radial coordinate (distance) to the center

  • phi (array of size n or float) – angular coordinate

  • center (array of size (2)) – rotation point

Returns:

array of same size with coords [x,y]

Raises:

AttributeError, KeyError

shear_polar2cartesian(phi, gamma)[source]
Parameters:
  • phi – shear angle (radian)

  • gamma – shear strength

Returns:

shear components gamma1, gamma2

shear_cartesian2polar(gamma1, gamma2)[source]
Parameters:
  • gamma1 – cartesian shear component

  • gamma2 – cartesian shear component

Returns:

shear angle, shear strength

phi_q2_ellipticity(phi, q)[source]

Transforms orientation angle and axis ratio into complex ellipticity moduli e1, e2.

Parameters:
  • phi – angle of orientation (in radian)

  • q – axis ratio minor axis / major axis

Returns:

eccentricities e1 and e2 in complex ellipticity moduli

ellipticity2phi_q(e1, e2)[source]

Transforms complex ellipticity moduli in orientation angle and axis ratio.

Parameters:
  • e1 – eccentricity in x-direction

  • e2 – eccentricity in xy-direction

Returns:

angle in radian, axis ratio (minor/major)

transform_e1e2_product_average(x, y, e1, e2, center_x, center_y)[source]

Maps the coordinates x, y with eccentricities e1 e2 into a new elliptical coordinate system such that R = sqrt(R_major * R_minor) :param x: x-coordinate :param y: y-coordinate :param e1: eccentricity :param e2: eccentricity :param center_x: center of distortion :param center_y: center of distortion :return: distorted coordinates x’, y’

Parameters:
  • x – x-coordinate

  • y – y-coordinate

  • e1 – eccentricity

  • e2 – eccentricity

  • center_x – center of distortion

  • center_y – center of distortion

Returns:

distorted coordinates x’, y’

transform_e1e2_square_average(x, y, e1, e2, center_x, center_y)[source]

Maps the coordinates x, y with eccentricities e1 e2 into a new elliptical coordinate system such that R = sqrt(R_major**2 + R_minor**2)

Parameters:
  • x – x-coordinate

  • y – y-coordinate

  • e1 – eccentricity

  • e2 – eccentricity

  • center_x – center of distortion

  • center_y – center of distortion

Returns:

distorted coordinates x’, y’

q2e(q)[source]

computes.

\[e = \equic \frac{1 - q^2}{1 + q^2}\]
Parameters:

q – axis ratio of minor to major axis

Returns:

ellipticity e

jaxtronomy.Util.util module

Module contents