import galsim as _galsim
import jax
import jax.numpy as jnp
import jax.scipy as jsp
from jax.tree_util import Partial as partial
from jax.tree_util import register_pytree_node_class
from jax_galsim.bessel import kv
from jax_galsim.core.draw import draw_by_kValue, draw_by_xValue
from jax_galsim.core.interpolate import akima_interp, akima_interp_coeffs
from jax_galsim.core.math import safe_sqrt
from jax_galsim.core.utils import cast_to_float, ensure_hashable, implements
from jax_galsim.gsobject import GSObject
from jax_galsim.random import UniformDeviate
@jax.jit
def _gamma(nu):
"""Gamma(nu) with care for integer nu in [0,5]"""
return jnp.select(
[nu == 0, nu == 1, nu == 2, nu == 3, nu == 4, nu == 5],
[jnp.inf, 1.0, 1.0, 2.0, 6.0, 24.0],
default=jsp.special.gamma(nu),
)
@jax.jit
def z2lz(z):
"""return z^2 * log(z)"""
return jnp.where(z <= 1e-40, 0.0, z * z * jnp.log(z))
@jax.jit
def f0(z):
"""K_0[z] with z -> 0 O(z^4)"""
z2 = z * z
z4 = z2 * z2
c0 = 0.11593151565841244881
c1 = 0.27898287891460311220
c2 = 0.025248929932162694513
return c0 + c1 * z2 + c2 * z4 - jnp.power(1.0 + 0.125 * z2, 2.0) * jnp.log(z)
@jax.jit
def f1(z):
"""z^1 K_1[z] with z -> 0 O(z^4)"""
z2 = z * z
z4 = z2 * z2
c0 = z2lz(z) # z^2 log(z)
c1 = 0.30796575782920622441
c2 = 0.08537071972865077805
return 1.0 - c1 * z2 - c2 * z4 + c0 * (0.5 + 0.0625 * z2)
@jax.jit
def f2(z):
"""z^2 K_2[z] with z -> 0 O(z^4)"""
c1 = 0.10824143945730155610
z2 = z * z
z4 = z2 * z2
c0 = z2lz(z) * z2 # z^4*log(z)
return 2.0 - 0.5 * z2 + c1 * z4 - 0.125 * c0
@jax.jit
def f3(z):
"""z^3 K_3[z] with z -> 0 O(z^4)"""
z2 = z * z
z4 = z2 * z2
return 8.0 - z2 + 0.125 * z4
@jax.jit
def f4(z):
"""z^4 K_4[z] with z -> 0 O(z^4)"""
z2 = z * z
z4 = z2 * z2
return 48.0 - 4 * z2 + 0.25 * z4
@jax.jit
def f5(z):
"""z^5 K_5[z] with z -> 0 O(z^4)"""
z2 = z * z
z4 = z2 * z2
return 384.0 - 24.0 * z2 + z4
@jax.jit
def fsmallz_nu(nu, z):
msk0 = nu == 0
msk1 = nu == 1
msk2 = nu == 2
msk3 = nu == 3
msk4 = nu == 4
nu_safe = jnp.where(
msk0 | msk1 | msk2 | msk3 | msk4,
nu + 1e-10,
nu,
)
def fnu(nu, z):
"""z^nu K_nu[z] with z -> 0 O(z^4) z > 0"""
z2 = z * z
z4 = z2 * z2
c1 = jnp.power(2.0, -6.0 - nu)
c2 = _gamma(-2.0 - nu)
c3 = _gamma(-2.0 + nu)
c4 = jnp.power(z, 2.0 * nu)
c5 = z4 * 8.0 * z2 * (2.0 + nu) + 32.0 * (1.0 + nu) * (2.0 + nu)
c6 = z2 * (16.0 + z2 - 8.0 * nu) * c3
return c1 * (c4 * c5 * c2 + jnp.power(4.0, nu) * (c6 + 32.0 * _gamma(nu)))
return jnp.select(
[msk0, msk1, msk2, msk3, msk4],
[f0(z), f1(z), f2(z), f3(z), f4(z)],
default=fnu(nu_safe, z),
)
@jax.jit
def _fz_nu(nu, z):
"""z^nu K_nu[z] with z > 0"""
return jnp.power(z, nu) * kv(nu, z)
@jax.jit
def fluxfractionFunc(z, nu, alpha):
"""1 - z^(nu+1) K_{nu+1}(z) / (2^nu Gamma(nu+1)) - alpha"""
return 1.0 - _fz_nu(nu + 1.0, z) / (jnp.power(2.0, nu) * _gamma(nu + 1.0)) - alpha
@jax.jit
def reducedfluxfractionFunc(z, nu, norm):
"""(1 - z^(nu+1) K_{nu+1}(z) / (2^nu Gamma(nu+1)))/norm"""
return fluxfractionFunc(z, nu, alpha=0.0) / norm
# code here is from JAX source for testing custom_root
# used under license:
# Copyright 2022 The JAX Authors.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# https://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
def _binary_search(func, x0, low=0.0, high=40.0):
del x0 # unused
def cond(state):
low, high = state
midpoint = 0.5 * (low + high)
return (low < midpoint) & (midpoint < high)
def body(state):
low, high = state
midpoint = 0.5 * (low + high)
update_upper = func(midpoint) > 0
low = jnp.where(update_upper, low, midpoint)
high = jnp.where(update_upper, midpoint, high)
return (low, high)
solution, _ = jax.lax.while_loop(cond, body, (low, high))
return solution
# end of code from jax source
@jax.jit
def calculateFluxRadius(alpha, nu, zmin=0, zmax=40):
"""Return radius R enclosing flux fraction alpha in unit of the scale radius r0
Method: Solve F(R/r0=z)/Flux - alpha = 0 using bisection algorithm
F(R)/F = int( 1/(2^nu Gamma(nu+1)) (r/r0)^(nu+1) K_nu(r/r0) dr/r0; r=0..R) = alpha
=>
z=R/r0 such that
1 - z^(nu+1) K_{nu+1}(z) / (2^nu Gamma(nu+1)) = alpha
Typical use cases:
o alpha = 1/2 => R = Half-Light-Radius,
o alpha = 1 - folding-thresold => R used for stepk computation
nu: the Spergel index
nb. it is supposed that nu is in [-0.85, 4.0] checked in the Spergel class init
"""
return jax.lax.custom_root(
partial(fluxfractionFunc, nu=nu, alpha=alpha),
20.0,
partial(_binary_search, low=zmin, high=zmax),
lambda f, y: y / f(1.0),
)
def _spergel_hlr_pade(x):
"""A Pseudo-Pade approximation for the HLR of the Spergel profile as a function of nu.
See dev/notebooks/spergel_hlr_flux_radius_approx.ipynb for code to generate this routine.
"""
# fmt: off
pm = 1.2571513771129166 + x * (
3.7059053890269102 + x * (
2.8577090425861944 + x * (
-0.30570486567039273 + x * (
0.6589831675940833 + x * (
3.375577680133867 + x * (
2.8143565844741403 + x * (
0.9292378858457211 + x * (
0.12096941981286179 + x * (
0.004206502758293099
)
)
)
)
)
)
)
)
)
qm = 1.0 + x * (
2.1939178810491837 + x * (
0.8281034080784796 + x * (
-0.5163329765186994 + x * (
0.9164871490929886 + x * (
1.8988551389326231 + x * (
1.042688817291684 + x * (
0.22580140592548198 + x * (
0.01681923980317362 + x * (
0.00018168506955933716
)
)
)
)
)
)
)
)
)
# fmt: on
return pm / qm
@jax.jit
def _spergel_hlr_binary_search_plus_pade_init(nu):
"""Return radius R enclosing flux fraction 0.5 in unit of the scale radius r0"""
z = _spergel_hlr_pade(nu)
# this parameter sets the window around the initial guess from the pade approximation
# it should be bigger than the maximum absolute error of the approximation
eps = 1e-6
def _hlr_bs(f, x0):
return _binary_search(f, x0, low=x0 - eps, high=x0 + eps)
return jax.lax.custom_root(
partial(fluxfractionFunc, nu=nu, alpha=0.5),
z,
_hlr_bs,
lambda f, y: y / f(1.0),
)
LAX_SPERGEL_DESCRIPTION = r"""
The fully normalized Spergel profile (used in both standard GalSim and JAX-GalSim) is
.. math::
I(r) = flux \times \left(2\pi 2^\nu \Gamma(1+\nu) r_0^2\right)^{-1} \times \left(\frac{r}{r_0}\right)^\nu K_\nu\left(\frac{r}{r_0}\right)
with the following Fourier expression
.. math::
\hat{I}(k) = flux / (1 + (k r_0)^2)^{1+\nu}
where :math:`r_0` is the ``scale_radius``, and :math:`\nu` mandatory to be in [-0.85,4.0]
The JAX-GalSim implementation does not support autodiff with respect to :math:`\nu` for
real-space evaluations.
When the profile is initialized using the half-light radius, if :math:`\nu` is a Python float,
then GalSim is used to convert to the scale radius :math:`r_0`. Otherwise, JAX-GalSim does the
conversion on-the-fly using a combination of an approximate guess and a binary search. Similar
logic is used for the ``calculateFluxRadius`` method.
"""
[docs]
@implements(_galsim.Spergel, lax_description=LAX_SPERGEL_DESCRIPTION)
@register_pytree_node_class
class Spergel(GSObject):
_has_hard_edges = False
_is_axisymmetric = True
_is_analytic_x = True
_is_analytic_k = True
_minimum_nu = -0.85
_maximum_nu = 4.0
def __init__(
self,
nu,
scale_radius=None,
half_light_radius=None,
flux=1.0,
gsparams=None,
):
nu = cast_to_float(nu)
# Parse the radius options
if half_light_radius is not None:
if scale_radius is not None:
raise _galsim.GalSimIncompatibleValuesError(
"Only one of scale_radius, half_light_radius may be specified",
half_light_radius=half_light_radius,
scale_radius=scale_radius,
)
else:
# for python floats, we can use galsim on the CPU-side to do this
# quickly as long as we ensure it is done at compile time.
if isinstance(nu, float):
with jax.ensure_compile_time_eval():
hlr = _galsim.Spergel(nu, scale_radius=1).half_light_radius
else:
hlr = _spergel_hlr_binary_search_plus_pade_init(nu)
super().__init__(
nu=nu,
scale_radius=half_light_radius / hlr,
flux=flux,
gsparams=gsparams,
)
elif scale_radius is None:
raise _galsim.GalSimIncompatibleValuesError(
"One of scale_radius, half_light_radius must be specified",
half_light_radius=half_light_radius,
scale_radius=scale_radius,
)
else:
super().__init__(
nu=nu,
scale_radius=scale_radius,
flux=flux,
gsparams=gsparams,
)
@property
@implements(_galsim.spergel.Spergel.nu)
def nu(self):
return self._params["nu"]
@property
@implements(_galsim.spergel.Spergel.scale_radius)
def scale_radius(self):
return self.params["scale_radius"]
@property
def _r0(self):
return self.scale_radius
@property
def _inv_r0(self):
return 1.0 / self._r0
@property
def _r0_sq(self):
return self._r0 * self._r0
@property
def _inv_r0_sq(self):
return self._inv_r0 * self._inv_r0
@property
@implements(_galsim.spergel.Spergel.half_light_radius)
def half_light_radius(self):
# for python floats, we can use galsim on the CPU-side to do this
# quickly as long as we ensure it is done at compile time.
if isinstance(self.nu, float):
with jax.ensure_compile_time_eval():
hlr = _galsim.Spergel(self.nu, scale_radius=1).half_light_radius
else:
hlr = _spergel_hlr_binary_search_plus_pade_init(self.nu)
return self._r0 * hlr
@property
def _shootxnorm(self):
"""Normalization for photon shooting"""
return 1.0 / (2.0 * jnp.pi * jnp.power(2.0, self.nu) * _gamma(self.nu + 1.0))
@property
def _xnorm(self):
"""Normalization of xValue"""
return self._shootxnorm * self.flux * self._inv_r0_sq
@property
def _xnorm0(self):
"""return z^nu K_nu(z) for z=0"""
return jax.lax.select(
self.nu > 0, _gamma(self.nu) * jnp.power(2.0, self.nu - 1.0), jnp.inf
)
[docs]
@implements(_galsim.spergel.Spergel.calculateFluxRadius)
def calculateFluxRadius(self, f):
f = cast_to_float(f)
# for python floats, we can use galsim on the CPU-side to do this
# quickly as long as we ensure it is done at compile time.
if isinstance(self.nu, float) and isinstance(f, float):
with jax.ensure_compile_time_eval():
fac = _galsim.Spergel(self.nu, scale_radius=1).calculateFluxRadius(f)
else:
fac = calculateFluxRadius(f, self.nu)
return self._r0 * fac
[docs]
@implements(_galsim.spergel.Spergel.calculateIntegratedFlux)
def calculateIntegratedFlux(self, r):
return fluxfractionFunc(r / self._r0, self.nu, 0.0)
def __hash__(self):
return hash(
(
"galsim.Spergel",
ensure_hashable(self.nu),
ensure_hashable(self.scale_radius),
ensure_hashable(self.flux),
self.gsparams,
)
)
def __repr__(self):
return "galsim.Spergel(nu=%r, scale_radius=%r, flux=%r, gsparams=%r)" % (
ensure_hashable(self.nu),
ensure_hashable(self.scale_radius),
ensure_hashable(self.flux),
self.gsparams,
)
def __str__(self):
s = "galsim.Spergel(nu=%s, half_light_radius=%s" % (
ensure_hashable(self.nu),
ensure_hashable(self.half_light_radius),
)
if self.flux != 1.0:
s += ", flux=%s" % (ensure_hashable(self.flux),)
s += ")"
return s
@property
def _maxk(self):
"""(1+ (k r0)^2)^(-1-nu) = maxk_threshold"""
res = jnp.power(self.gsparams.maxk_threshold, -1.0 / (1.0 + self.nu)) - 1.0
return jnp.sqrt(res) / self._r0
@property
def _stepk(self):
R = calculateFluxRadius(1.0 - self.gsparams.folding_threshold, self.nu)
R *= self._r0
# Go to at least 5*hlr
R = jnp.maximum(R, self.gsparams.stepk_minimum_hlr * self.half_light_radius)
return jnp.pi / R
@property
def _max_sb(self):
# from SBSpergelImpl.h
return jnp.abs(self._xnorm) * self._xnorm0
@staticmethod
@jax.jit
def _xValue_exact_func(nu, r, xnorm):
return xnorm * _fz_nu(nu, r)
@staticmethod
@jax.jit
def _xValue_asymp_func(nu, r, xnorm):
return xnorm * jnp.power(r, nu) * jnp.exp(-r) * jnp.sqrt(jnp.pi / 2 / r)
@staticmethod
@jax.jit
def _xValue_smallz_func(nu, r, xnorm):
return xnorm * fsmallz_nu(nu, r)
def _xValue_interp_coeffs(self):
# MRB: this number of points gets the tests to pass
# I did not investigate further.
n_pts = 2000
r_min = jnp.minimum(jnp.pi / self.maxk, 1e-6)
r_max = jnp.pi / self.stepk
r = jnp.logspace(jnp.log10(r_min), jnp.log10(r_max), n_pts)
nu = jax.lax.stop_gradient(self.nu)
vals = self._xValue_exact_func(
nu,
r / self._r0,
jnp.abs(self._xnorm),
)
# slope to match the interpolant onto an asymptotic expansion of kv
# that is kv(x) ~ sqrt(pi/2/x) * exp(-x) * (1 + slp/x)
xval = r[-1] / self._r0
aval = self._xValue_asymp_func(nu, xval, jnp.abs(self._xnorm))
slp_asymp = (vals[-1] / aval - 1) * xval
# slope to match the interpolant onto a taylor expansion to z^4
# via kv(x) ~ kv_smallz(x) * (1 + slp * x**4)
xval = r[0] / self._r0
aval = self._xValue_smallz_func(nu, xval, jnp.abs(self._xnorm))
slp_smallz = (vals[0] / aval - 1) / jnp.power(xval, 4)
return (
r,
jnp.log(vals),
akima_interp_coeffs(jnp.log(r), jnp.log(vals)),
slp_asymp,
slp_smallz,
)
@jax.jit
def _xValue(self, pos):
# we cannot compute gradients with respect to nu
nu = jax.lax.stop_gradient(self.nu)
r = safe_sqrt(pos.x**2 + pos.y**2)
# we work with a 1D array and reshape it back at the end
out_shape = jnp.shape(r)
r = jnp.atleast_1d(r).ravel()
# the computation here uses 4 parts
# - a value at r = 0
# - a taylor expansion at small r
# - an interpolant
# - an asymptotic expansion at large r
# the interpolant is matched onto the taylor and asymptotic via two slopes
r_, vals_, coeffs, slp_asymp, slp_smallz = self._xValue_interp_coeffs()
# define masks for each range
msk_nz = r > 0
r_msk_nz = jnp.where(msk_nz, r, r_[0])
r_msk_nz_inv_r0 = r_msk_nz * self._inv_r0
msk_asymp = r > r_[-1]
msk_smallz = r < r_[0]
msk_interp = (~msk_smallz) & (~msk_asymp)
msk_smallz = msk_smallz & msk_nz
# compute values for each range
res_z = self._xnorm0 * self._xnorm
res_smallz = self._xValue_smallz_func(nu, r_msk_nz_inv_r0, self._xnorm) * (
1.0 + slp_smallz * jnp.power(r_msk_nz_inv_r0, 4)
)
res_interp = jnp.exp(
akima_interp(
jnp.log(r_msk_nz), jnp.log(r_), vals_, coeffs, fixed_spacing=True
)
) * jnp.sign(self._xnorm)
res_asymp = self._xValue_asymp_func(
nu,
r_msk_nz_inv_r0,
self._xnorm,
) * (1.0 + slp_asymp / r_msk_nz_inv_r0)
# pick the right value
res = jnp.select(
[~msk_nz, msk_smallz, msk_interp, msk_asymp],
[res_z, res_smallz, res_interp, res_asymp],
)
# reshape to final output
return res.reshape(out_shape)
@jax.jit
def _kValue(self, kpos):
ksq = (kpos.x**2 + kpos.y**2) * self._r0_sq
return self.flux * jnp.power(1.0 + ksq, -1.0 - self.nu)
def _drawReal(self, image, jac=None, offset=(0.0, 0.0), flux_scaling=1.0):
_jac = jnp.eye(2) if jac is None else jac
return draw_by_xValue(self, image, _jac, jnp.asarray(offset), flux_scaling)
def _drawKImage(self, image, jac=None):
_jac = jnp.eye(2) if jac is None else jac
return draw_by_kValue(self, image, _jac)
[docs]
@implements(_galsim.Spergel.withFlux)
def withFlux(self, flux):
return Spergel(
nu=self.nu,
scale_radius=self.scale_radius,
flux=flux,
gsparams=self.gsparams,
)
@property
def _shoot_pos_cdf(self):
zmax = calculateFluxRadius(
1.0 - self.gsparams.shoot_accuracy, self.nu, zmax=30.0
)
flux_max = fluxfractionFunc(zmax, self.nu, alpha=0.0)
preducedfluxfractionFunc = partial(
reducedfluxfractionFunc, nu=self.nu, norm=flux_max
)
z_cdf = jnp.linspace(0, zmax, 10_000)
cdf = preducedfluxfractionFunc(z_cdf)
return z_cdf, cdf
def _shoot_pos(self, u):
# shoot r in case of nu>0
z_cdf, cdf = self._shoot_pos_cdf
z = jnp.interp(u, cdf, z_cdf) # linear inversion of the CDF
r = z * self._r0
return r
@property
def _shoot_neg_cdf(self):
# comment:
# In the Galsim code the profile below rmin is linearized such that
# call zmin = rmin/r0 such that
# Int_0^zmin 2pi u x I(u) du = shoot_accuracy
# Then let (a,b) such that
# 1) Int_0^zmin 2pi u x (a + b u) du = shoot_accuracy
# 2) a + b zmin = zmin^nu K_nu(zmin)
# Now, noticing that
# I(z) = z^nu K_nu(z) / (2pi 2^nu Gamma(nu+1)) = z^nu K_nu(z)/(2 pi Nnu)
# there is a problem with eq. 1 as we would have expected
# 1b) Int_0^zmin 2pi u x (a + b u)/(2 pi Nnu) du = shoot_accuracy
# so the corrFact is there to signal the changement in this implementation
zmax = calculateFluxRadius(
1.0 - self.gsparams.shoot_accuracy, self.nu, zmax=30.0
)
flux_target = self.gsparams.shoot_accuracy
shoot_rmin = calculateFluxRadius(flux_target, self.nu)
knur = _fz_nu(self.nu, shoot_rmin)
corrFact = self._shootxnorm # this is the correct normalisation
b = knur - flux_target / (jnp.pi * shoot_rmin * shoot_rmin * corrFact)
b = 3.0 * b / shoot_rmin
a = knur - shoot_rmin * b
def cumulflux(z, a, b, zmin, nu, norm=1.0):
flux_min = a / 3.0 * zmin * zmin * zmin + b / 2.0 * zmin * zmin
c1 = _fz_nu(nu + 1.0, zmin)
res = jnp.where(
z <= zmin,
a / 3.0 * z * z * z + b / 2.0 * z * z,
flux_min + c1 - _fz_nu(nu + 1.0, z),
)
return res / norm
flux_max = cumulflux(zmax, a, b, shoot_rmin, self.nu)
preducedfluxfractionFunc = partial(
cumulflux, a=a, b=b, zmin=shoot_rmin, nu=self.nu, norm=flux_max
)
z_cdf = jnp.linspace(0, zmax, 10_000)
cdf = preducedfluxfractionFunc(z_cdf)
return z_cdf, cdf
def _shoot_neg(self, u):
# shoot r in case of nu<=0
z_cdf, cdf = self._shoot_neg_cdf
z = jnp.interp(u, cdf, z_cdf) # linear inversion of the CDF
r = z * self._r0
return r
@implements(_galsim.Spergel._shoot)
def _shoot(self, photons, rng):
ud = UniformDeviate(rng)
u = ud.generate(photons.x)
r = jax.lax.select(self.nu > 0, self._shoot_pos(u), self._shoot_neg(u))
ang = ud.generate(photons.x) * 2.0 * jnp.pi
photons.x = r * jnp.cos(ang)
photons.y = r * jnp.sin(ang)
photons.flux = self.flux / photons.size()