Spergel

CHANGELOG.md

@@ -24,7 +24,7 @@
     * `makePhot`
     * `drawPhot`
   * Added implementation of simple light profiles:
-    * `Gaussian`, `Exponential`, `Pixel`, `Box`, `Moffat`, `DeltaFunction`
+    * `Gaussian`, `Exponential`, `Pixel`, `Box`, `Moffat`, `Spergel`, `DeltaFunction`
   * Added implementation of simple WCS:
     * `PixelScale`, `OffsetWCS`, `JacobianWCS`, `AffineTransform`, `ShearWCS`, `OffsetShearWCS`, `GSFitsWCS`, `FitsWCS`, `TanWCS`
   * Added automated suite of tests using the reference GalSim and LSSTDESC-Coord test suites

jax_galsim/__init__.py

@@ -53,6 +53,7 @@
 from .box import Box, Pixel
 from .gsobject import GSObject
 from .moffat import Moffat
+from .spergel import Spergel
 from .sum import Add, Sum
 from .transform import Transform, Transformation
 from .convolve import Convolve, Convolution, Deconvolution, Deconvolve

jax_galsim/bessel.py

@@ -2,6 +2,7 @@
 import jax
 import jax.numpy as jnp
 from jax._src.numpy.util import _wraps
+import tensorflow_probability as tfp
 
 
 # the code here for Si, f, g and _si_small_pade is taken from galsim/src/math/Sinc.cpp
@@ -101,3 +102,15 @@ def si(x):
         - _g_pade(x, x2) * jnp.sin(x),
         _si_small_pade(x, x2),
     )
+
+
+@jax.jit
+def kv(nu, x):
+    """Modified Bessel 2nd kind"""
+    return tfp.substrates.jax.math.bessel_kve(nu * 1.0, x) / jnp.exp(jnp.abs(x))
+
+
+@jax.jit
+def gamma(x):
+    """Gamma(x)"""
+    return jnp.exp(jax.lax.lgamma(x * 1.0))

jax_galsim/moffat.py

@@ -7,6 +7,7 @@
 from jax.tree_util import register_pytree_node_class
 
 from jax_galsim.core.bessel import j0
+from jax_galsim.bessel import kv
 from jax_galsim.core.draw import draw_by_kValue, draw_by_xValue
 from jax_galsim.core.integrate import ClenshawCurtisQuad, quad_integral
 from jax_galsim.core.utils import bisect_for_root, ensure_hashable
@@ -18,7 +19,7 @@
 @jax.jit
 def _Knu(nu, x):
     """Modified Bessel 2nd kind for Untruncated Moffat"""
-    return tfp.substrates.jax.math.bessel_kve(nu * 1.0, x) / jnp.exp(jnp.abs(x))
+    return kv(nu, x)
 
 
 @jax.jit