Core Utilities¶

Please note that items from the jax_galsim.core are internal APIs and should not be relied upon by external code.

Math¶

jax_galsim.core.math.safe_sqrt(x)[source]¶: Numerically safe sqrt operation with zero derivative at zero.

jax_galsim.core.integrate.abs_weights(n)[source]¶

Parameters:: n (int)

class jax_galsim.core.integrate.ClenshawCurtisQuad(order, absc, absw, errw)[source]¶

Bases: NamedTuple

Parameters:

order (int)
absc (Array)
absw (Array)
errw (Array)

order: int¶: Alias for field number 0

absc: Array¶: Alias for field number 1

absw: Array¶: Alias for field number 2

errw: Array¶: Alias for field number 3

classmethod init(order)[source]¶

Parameters:: order (int)

static compute_weights(order)[source]¶

Parameters:: order (int)

static rescale_weights(absc, absw, *, interval_in=(-1, 1), interval_out=(0, 1))[source]¶

Parameters:

absc (Array)
absw (Array)
interval_in (Tuple[float, float])
interval_out (Tuple[float, float])

jax_galsim.core.integrate.quad_integral(f, a, b, quad)[source]¶

Parameters:: quad (ClenshawCurtisQuad)

jax_galsim.core.interpolate.akima_interp_coeffs(x, y, use_jax=True)[source]¶

Compute the interpolation coefficients for an Akima cubic spline.

An Akima cubic spline is a piecewise C(1) cubic polynomial that interpolates a set of points (x, y). Unlike a more traditional cubic spline, the Akima spline can be computed without solving a linear system of equations. However, the Akima spline does not have continuous second derivatives at the interpolation points.

See https://en.wikipedia.org/wiki/Akima_spline and Akima (1970), “A new method of interpolation and smooth curve fitting based on local procedures”, Journal of the ACM. 17: 589-602 for a description of the technique.

Parameters:

x – The x-coordinates of the data points. These must be sorted into increasing order and cannot contain any duplicates.
y – The y-coordinates of the data points.
use_jax – Whether to use JAX for computation. If False, coefficients are computed using NumPy on the host device, which can be useful when embedded inside JAX code with JIT applied to pre-compute the coefficients. [default: True]

Returns:

A tuple of arrays (a, b, c, d) where each array has shape (N-1,) and contains the coefficients for the cubic polynomial that interpolates the data points between x[i] and x[i+1].

jax_galsim.core.interpolate.akima_interp(x, xp, yp, coeffs, fixed_spacing=False)[source]¶

Compute the values of an Akima cubic spline at a set of points given the interpolation coefficients.

Parameters:

x – The x-coordinates of the points where the interpolation is computed.
xp – The x-coordinates of the data points. These must be sorted into increasing order and cannot contain any duplicates.
yp – The y-coordinates of the data points. Not used currently.
coeffs – The interpolation coefficients returned by akima_interp_coeffs.
fixed_spacing – Whether the data points are evenly spaced. If True, a faster technique is used to find the index i such that xp[i] <= x < xp[i+1]. [default: False]

Returns:

The values of the Akima cubic spline at the points x.

Drawing¶

jax_galsim.core.draw.draw_by_xValue(gsobject, image, jacobian=Array([[1., 0.], [0., 1.]], dtype=float32), offset=Array([0., 0.], dtype=float32), flux_scaling=1.0)[source]¶: Utility function to draw a real-space GSObject into an Image.

jax_galsim.core.draw.draw_by_kValue(gsobject, image, jacobian=Array([[1., 0.], [0., 1.]], dtype=float32))[source]¶

jax_galsim.core.draw.apply_kImage_phases(offset, image, jacobian=Array([[1., 0.], [0., 1.]], dtype=float32))[source]¶

jax_galsim.core.draw.calculate_mean_n_photons(flux, flux_per_photon, max_sb)[source]¶

Calculate the mean number of photons to shoot for photon shooting.

This routine can be used to group objects together by the typical number of photons they will shoot when drawing objects in bulk.

Parameters:

flux – The flux of the GSObject (e.g., obj.flux).
flux_per_photon – The flux per photon (e.g., obj._flux_per_photon).
max_sb – The maximum surface brightness of the object (e.g., obj.max_sb).

Returns:

The number of photons.

jax_galsim.core.draw.calculate_n_photons(flux, flux_per_photon, max_sb, n_photons=0, rng=None, max_extra_noise=0.0, poisson_flux=True)[source]¶

Calculate the number of photons to shoot for an object when photon shooting according to the code in galsim.GSObject._calculate_n_photons. See the notes section below for more details.

Parameters:

flux – The flux of the GSObject (e.g., obj.flux).
flux_per_photon – The flux per photon (e.g., obj._flux_per_photon).
max_sb – The maximum surface brightness of the object (e.g., obj.max_sb).
n_photons – If provided, the number of photons to use for photon shooting. If not provided (i.e. n_photons = 0), use as many photons as necessary to result in an image with the correct Poisson shot noise for the object’s flux. For positive definite profiles, this is equivalent to n_photons = flux. However, some profiles need more than this because some of the shot photons are negative (usually due to interpolants). [default: 0]
rng – If provided, a random number generator to use for photon shooting, which may be any kind of BaseDeviate object. If rng is None, one will be automatically created, using the time as a seed. [default: None]
max_extra_noise – If provided, the allowed extra noise in each pixel when photon shooting. This is only relevant if n_photons=0, so the number of photons is being automatically calculated. In that case, if the image noise is dominated by the sky background, then you can get away with using fewer shot photons than the full n_photons = flux. Essentially each shot photon can have a flux > 1, which increases the noise in each pixel. The max_extra_noise parameter specifies how much extra noise per pixel is allowed because of this approximation. A typical value for this might be max_extra_noise = sky_level / 100 where sky_level is the flux per pixel due to the sky. Note that this uses a “variance” definition of noise, not a “sigma” definition. [default: 0.]
poisson_flux – Whether to allow total object flux scaling to vary according to Poisson statistics for n_photons samples when photon shooting. [default: True, unless n_photons is given, in which case the default is False]

Returns:

A tuple of (n_photons, g, rng) where n_photons is the number of photons, g is the flux ratio, and rng is the final random number generator used.

Notes

It is easiest to look at the original code from GSObject._calculate_nphotons to understand what this function does:

# Calculate how many photons to shoot and what flux_ratio (called g) each one should
# have in order to produce an image with the right S/N and total flux.
#
# This routine is normally called by drawPhot.
#
# Returns:
#     n_photons, g

# For profiles that are positive definite, then N = flux. Easy.
#
# However, some profiles shoot some of their photons with negative flux. This means that
# we need a few more photons to get the right S/N = sqrt(flux). Take eta to be the
# fraction of shot photons that have negative flux.
#
# S^2 = (N+ - N-)^2 = (N+ + N- - 2N-)^2 = (Ntot - 2N-)^2 = Ntot^2(1 - 2 eta)^2
# N^2 = Var(S) = (N+ + N-) = Ntot
#
# So flux = (S/N)^2 = Ntot (1-2eta)^2
# Ntot = flux / (1-2eta)^2
#
# However, if each photon has a flux of 1, then S = (1-2eta) Ntot = flux / (1-2eta).
# So in fact, each photon needs to carry a flux of g = 1-2eta to get the right
# total flux.
#
# That's all the easy case. The trickier case is when we are sky-background dominated.
# Then we can usually get away with fewer shot photons than the above.  In particular,
# if the noise from the photon shooting is much less than the sky noise, then we can
# use fewer shot photons and essentially have each photon have a flux > 1. This is ok
# as long as the additional noise due to this approximation is "much less than" the
# noise we'll be adding to the image for the sky noise.
#
# Let's still have Ntot photons, but now each with a flux of g. And let's look at the
# noise we get in the brightest pixel that has a nominal total flux of Imax.
#
# The number of photons hitting this pixel will be Imax/flux * Ntot.
# The variance of this number is the same thing (Poisson counting).
# So the noise in that pixel is:
#
# N^2 = Imax/flux * Ntot * g^2
#
# And the signal in that pixel will be:
#
# S = Imax/flux * (N+ - N-) * g which has to equal Imax, so
# g = flux / Ntot(1-2eta)
# N^2 = Imax/Ntot * flux / (1-2eta)^2
#
# As expected, we see that lowering Ntot will increase the noise in that (and every
# other) pixel.
# The input max_extra_noise parameter is the maximum value of spurious noise we want
# to allow.
#
# So setting N^2 = Imax + nu, we get
#
# Ntot = flux / (1-2eta)^2 / (1 + nu/Imax)
# g = (1 - 2eta) * (1 + nu/Imax)
#
# Returns the total flux placed inside the image bounds by photon shooting.

flux = self.flux
if flux == 0.0:
    return 0, 1.0

# The _flux_per_photon property is (1-2eta)
# This factor will already be accounted for by the shoot function, so don't include
# that as part of our scaling here.  There may be other adjustments though, so g=1 here.
eta_factor = self._flux_per_photon
mod_flux = flux / (eta_factor * eta_factor)
g = 1.0

# If requested, let the target flux value vary as a Poisson deviate
if poisson_flux:
    # If we have both positive and negative photons, then the mix of these
    # already gives us some variation in the flux value from the variance
    # of how many are positive and how many are negative.
    # The number of negative photons varies as a binomial distribution.
    # <F-> = eta * Ntot * g
    # <F+> = (1-eta) * Ntot * g
    # <F+ - F-> = (1-2eta) * Ntot * g = flux
    # Var(F-) = eta * (1-eta) * Ntot * g^2
    # F+ = Ntot * g - F- is not an independent variable, so
    # Var(F+ - F-) = Var(Ntot*g - 2*F-)
    #              = 4 * Var(F-)
    #              = 4 * eta * (1-eta) * Ntot * g^2
    #              = 4 * eta * (1-eta) * flux
    # We want the variance to be equal to flux, so we need an extra:
    # delta Var = (1 - 4*eta + 4*eta^2) * flux
    #           = (1-2eta)^2 * flux
    absflux = abs(flux)
    mean = eta_factor * eta_factor * absflux
    pd = PoissonDeviate(rng, mean)
    pd_val = pd() - mean + absflux
    ratio = pd_val / absflux
    g *= ratio
    mod_flux *= ratio

if n_photons == 0.0:
    n_photons = abs(mod_flux)
    if max_extra_noise > 0.0:
        gfactor = 1.0 + max_extra_noise / abs(self.max_sb)
        n_photons /= gfactor
        g *= gfactor

# Make n_photons an integer.
iN = int(n_photons + 0.5)

return iN, g

Utilities¶

jax_galsim.core.utils.cast_numpy_array_to_native_byte_order(arr)[source]¶: Cast an array to native byte order.

jax_galsim.core.utils.has_tracers(x)[source]¶: Return True if the input item is a JAX tracer or object, False otherwise.

jax_galsim.core.utils.compute_major_minor_from_jacobian(jac)[source]¶

jax_galsim.core.utils.cast_to_python_float(x)[source]¶: Cast the input to a python float. Works on python floats and jax arrays. For jax arrays it always takes the first element after a call to .ravel()

jax_galsim.core.utils.cast_to_python_int(x)[source]¶: Cast the input to a python int. Works on python ints and jax arrays. For jax arrays it always takes the first element after a call to .ravel()

jax_galsim.core.utils.cast_to_float(x)[source]¶: Cast the input to a float. Works on python floats and jax arrays.

jax_galsim.core.utils.cast_to_int(x)[source]¶: Cast the input to an int. Works on python floats/ints and jax arrays.

jax_galsim.core.utils.is_equal_with_arrays(x, y)[source]¶: Return True if the data is equal, False otherwise. Handles jax.Array types.

jax_galsim.core.utils.ensure_hashable(v)[source]¶

Ensure that the input is hashable. If it is a jax array, convert it to a possibly nested tuple or python scalar.

All NaNs are converted to numpy.nan to get consistent hashing.

jax_galsim.core.utils.bisect_for_root(func, low, high, niter=75)[source]¶

class jax_galsim.core.utils.ParsedDoc(docstr='', signature='', summary='', front_matter='', sections={})[source]¶

Bases: NamedTuple

docstr: full docstring signature: signature from docstring. summary: summary from docstring. front_matter: front matter before sections. sections: dictionary of section titles to section content.

Parameters:

docstr (str)
signature (str)
summary (str)
front_matter (str)
sections (dict[str, str])

docstr: str¶: Alias for field number 0

signature: str¶: Alias for field number 1

summary: str¶: Alias for field number 2

front_matter: str¶: Alias for field number 3

sections: dict[str, str]¶: Alias for field number 4

jax_galsim.core.utils.implements(original_fun, lax_description='', module=None)[source]¶

Decorator for JAX functions which implement a specified GalSim function.

This mainly contains logic to copy and modify the docstring of the original function. In particular, if update_doc is True, parameters listed in the original function that are not supported by the decorated function will be removed from the docstring. For this reason, it is important that parameter names match those in the original GalSim function.

Parameters:

original_fun – The original function being implemented
lax_description – A string description that will be added to the beginning of the docstring.
module – An optional string specifying the module from which the original function is imported. This is useful for objects, where the module cannot be determined from the original function itself.