fix: do not use laxy properties or workspaces

Author: beckermr

jax_galsim/exponential.py

@@ -1,12 +1,14 @@
+from functools import lru_cache
+
 import galsim as _galsim
 import jax.numpy as jnp
+import numpy as np
 from jax.tree_util import register_pytree_node_class
 
 from jax_galsim.core.draw import draw_by_kValue, draw_by_xValue
 from jax_galsim.core.utils import ensure_hashable, implements
 from jax_galsim.gsobject import GSObject
 from jax_galsim.random import UniformDeviate
-from jax_galsim.utilities import lazy_property
 
 
 @implements(_galsim.Exponential)
@@ -147,50 +149,14 @@ def withFlux(self, flux):
             scale_radius=self.scale_radius, flux=flux, gsparams=self.gsparams
         )
 
-    @lazy_property
-    def _shoot_cdf(self):
-        # Comments on the math here:
-        #
-        # We are looking to draw from a distribution that is r * exp(-r).
-        # This distribution is the radial PDF of an Exponential profile.
-        # The factor of r comes from the area element r * dr.
-        #
-        # We can compute the CDF of this distribution analytically, but we cannot
-        # invert the CDF in closed form. Thus we invert it numerically using a table.
-        #
-        # One final detail is that we want the inversion to be accurate and are using
-        # linear interpolation. Thus we use a change of variables r = -ln(1 - u)
-        # to make the CDF more linear and map it's domain to [0, 1) instead of [0, inf).
-        #
-        # Putting this all together, we get
-        #
-        #     r * exp(-r) dr = -ln(1-u) (1-u) dr/du du
-        #                    = -ln(1-u) (1-u)  * 1 / (1-u)
-        #                    = -ln(1-u)
-        #
-        # The new range of integration is u = 0 to u = 1. Thus the CDF is
-        #
-        #     CDF = -int_0^u ln(1-u') du'
-        #         =  u - (u - 1) ln(1 - u)
-        #
-        # The final detail is that galsim defines a shoot accuracy and draws photons
-        # between r = 0 and rmax = -log(shoot_accuracy). Thus we normalize the CDF to
-        # its value at umax = 1 - exp(-rmax) and then finally invert the CDF numerically.
-        _rmax = -jnp.log(self.gsparams.shoot_accuracy)
-        _umax = 1.0 - jnp.exp(-_rmax)
-        _u_cdf = jnp.linspace(0, _umax, 10000)
-        _cdf = _u_cdf - (_u_cdf - 1) * jnp.log(1 - _u_cdf)
-        _cdf /= _cdf[-1]
-        return _u_cdf, _cdf
-
     @implements(_galsim.Exponential._shoot)
     def _shoot(self, photons, rng):
         ud = UniformDeviate(rng)
 
         u = ud.generate(
             photons.x
         )  # this does not fill arrays like in galsim so is safe
-        _u_cdf, _cdf = self._shoot_cdf
+        _u_cdf, _cdf = _shoot_cdf(self.gsparams.shoot_accuracy)
         # this interpolation inverts the CDF
         u = jnp.interp(u, _cdf, _u_cdf)
         # this converts from u (see above) to r and scales by the actual size of
@@ -203,3 +169,42 @@ def _shoot(self, photons, rng):
         photons.x = r * jnp.cos(ang)
         photons.y = r * jnp.sin(ang)
         photons.flux = self.flux / photons.size()
+
+
+@lru_cache(maxsize=8)
+def _shoot_cdf(shoot_accuracy):
+    """This routine produces a CPU-side cache of the CDF that is embedded
+    into JIT-compiled code as needed."""
+    # Comments on the math here:
+    #
+    # We are looking to draw from a distribution that is r * exp(-r).
+    # This distribution is the radial PDF of an Exponential profile.
+    # The factor of r comes from the area element r * dr.
+    #
+    # We can compute the CDF of this distribution analytically, but we cannot
+    # invert the CDF in closed form. Thus we invert it numerically using a table.
+    #
+    # One final detail is that we want the inversion to be accurate and are using
+    # linear interpolation. Thus we use a change of variables r = -ln(1 - u)
+    # to make the CDF more linear and map it's domain to [0, 1) instead of [0, inf).
+    #
+    # Putting this all together, we get
+    #
+    #     r * exp(-r) dr = -ln(1-u) (1-u) dr/du du
+    #                    = -ln(1-u) (1-u)  * 1 / (1-u)
+    #                    = -ln(1-u)
+    #
+    # The new range of integration is u = 0 to u = 1. Thus the CDF is
+    #
+    #     CDF = -int_0^u ln(1-u') du'
+    #         =  u - (u - 1) ln(1 - u)
+    #
+    # The final detail is that galsim defines a shoot accuracy and draws photons
+    # between r = 0 and rmax = -log(shoot_accuracy). Thus we normalize the CDF to
+    # its value at umax = 1 - exp(-rmax) and then finally invert the CDF numerically.
+    _rmax = -np.log(shoot_accuracy)
+    _umax = 1.0 - np.exp(-_rmax)
+    _u_cdf = np.linspace(0, _umax, 10000)
+    _cdf = _u_cdf - (_u_cdf - 1) * np.log(1 - _u_cdf)
+    _cdf /= _cdf[-1]
+    return _u_cdf, _cdf

jax_galsim/gsobject.py

@@ -29,7 +29,6 @@ class GSObject:
     def __init__(self, *, gsparams=None, **params):
         self._params = params  # Dictionary containing all traced parameters
         self._gsparams = GSParams.check(gsparams)  # Non-traced static parameters
-        self._workspace = {}  # used by lazy_property
 
     def __getstate__(self):
         d = self.__dict__.copy()

jax_galsim/interpolant.py

@@ -20,7 +20,6 @@
 from jax_galsim.errors import GalSimError
 from jax_galsim.gsparams import GSParams
 from jax_galsim.random import UniformDeviate
-from jax_galsim.utilities import lazy_property
 
 
 @implements(_galsim.interpolant.Interpolant)
@@ -225,7 +224,7 @@ def urange(self):
             % self.__class__.__name__
         )
 
-    @lazy_property
+    # TODO: Work out CPU-side caching and pre-generation for this
     def _shoot_cdf(self):
         x = jnp.linspace(-self.xrange, self.xrange, 10000)
         px = jnp.abs(self._xval_noraise(jnp.abs(x)))
@@ -1389,7 +1388,7 @@ def _du(self):
             / self._n
         )
 
-    @lazy_property
+    @property
     def _umax(self):
         return _find_umax_lanczos(
             self._du,

jax_galsim/interpolatedimage.py

@@ -31,7 +31,7 @@
 from jax_galsim.position import PositionD
 from jax_galsim.random import UniformDeviate
 from jax_galsim.transform import Transformation
-from jax_galsim.utilities import convert_interpolant, lazy_property
+from jax_galsim.utilities import convert_interpolant
 from jax_galsim.wcs import BaseWCS, PixelScale
 
 # These keys are removed from the public API of
@@ -548,7 +548,7 @@ def k_interpolant(self):
         """The Fourier-space `Interpolant` for this profile."""
         return self._k_interpolant
 
-    @lazy_property
+    @property
     def image(self):
         """The underlying `Image` being interpolated."""
         return self._xim[self._image.bounds]
@@ -557,19 +557,19 @@ def image(self):
     def _flux(self):
         return self._image_flux
 
-    @lazy_property
+    @property
     def _centroid(self):
         x, y = self._pad_image.get_pixel_centers()
         tot = jnp.sum(self._pad_image.array)
         xpos = jnp.sum(x * self._pad_image.array) / tot
         ypos = jnp.sum(y * self._pad_image.array) / tot
         return PositionD(xpos, ypos)
 
-    @lazy_property
+    @property
     def _max_sb(self):
         return jnp.max(jnp.abs(self._pad_image.array))
 
-    @lazy_property
+    @property
     def _flux_ratio(self):
         if self._jax_children[1]["flux"] is None:
             flux = self._image_flux
@@ -585,11 +585,11 @@ def _flux_ratio(self):
         # this class
         return flux / self._image_flux
 
-    @lazy_property
+    @property
     def _image_flux(self):
         return jnp.sum(self._image.array, dtype=float)
 
-    @lazy_property
+    @property
     def _offset(self):
         # Figure out the offset to apply based on the original image (not the padded one).
         # We will apply this below in _sbp.
@@ -598,7 +598,7 @@ def _offset(self):
             self._image.bounds, offset, None, self._jax_aux_data["use_true_center"]
         )
 
-    @lazy_property
+    @property
     def _image(self):
         # Store the image as an attribute and make sure we don't change the original image
         # in anything we do here.  (e.g. set scale, etc.)
@@ -616,7 +616,7 @@ def _image(self):
 
         return image
 
-    @lazy_property
+    @property
     def _wcs(self):
         im_cen = (
             self._jax_children[0].true_center
@@ -634,15 +634,15 @@ def _wcs(self):
 
         return wcs.local(image_pos=im_cen)
 
-    @lazy_property
+    @property
     def _jac_arr(self):
         image = self._jax_children[0]
         im_cen = (
             image.true_center if self._jax_aux_data["use_true_center"] else image.center
         )
         return self._wcs.jacobian(image_pos=im_cen).getMatrix().ravel()
 
-    @lazy_property
+    @property
     def _xim(self):
         pad_factor = self._jax_aux_data["pad_factor"]
 
@@ -669,27 +669,27 @@ def _xim(self):
 
         return xim
 
-    @lazy_property
+    @property
     def _pad_image(self):
         # These next two allow for easy pickling/repring.  We don't need to serialize all the
         # zeros around the edge.  But we do need to keep any non-zero padding as a pad_image.
         xim = self._xim
         nz_bounds = self._image.bounds
         return xim[nz_bounds]
 
-    @lazy_property
+    @property
     def _kim(self):
         return self._xim.calculate_fft()
 
-    @lazy_property
+    @property
     def _maxk(self):
         if self._jax_aux_data["_force_maxk"]:
             _, minor = compute_major_minor_from_jacobian(self._jac_arr.reshape((2, 2)))
             return self._jax_aux_data["_force_maxk"] * minor
         else:
             return self._getMaxK(self._jax_aux_data["calculate_maxk"])
 
-    @lazy_property
+    @property
     def _stepk(self):
         if self._jax_aux_data["_force_stepk"]:
             _, minor = compute_major_minor_from_jacobian(self._jac_arr.reshape((2, 2)))
@@ -837,7 +837,7 @@ def _drawKImage(self, image, jac=None):
         # Return an image
         return Image(array=im, bounds=image.bounds, wcs=image.wcs, _check_bounds=False)
 
-    @lazy_property
+    @property
     def _pos_neg_fluxes(self):
         # record pos and neg fluxes now too
         pflux = jnp.sum(jnp.where(self._pad_image.array > 0, self._pad_image.array, 0))

jax_galsim/spergel.py

@@ -9,7 +9,6 @@
 from jax_galsim.core.utils import bisect_for_root, ensure_hashable, implements
 from jax_galsim.gsobject import GSObject
 from jax_galsim.random import UniformDeviate
-from jax_galsim.utilities import lazy_property
 
 
 @jax.jit
@@ -321,34 +320,34 @@ def scale_radius(self):
     def _r0(self):
         return self.scale_radius
 
-    @lazy_property
+    @property
     def _inv_r0(self):
         return 1.0 / self._r0
 
-    @lazy_property
+    @property
     def _r0_sq(self):
         return self._r0 * self._r0
 
-    @lazy_property
+    @property
     def _inv_r0_sq(self):
         return self._inv_r0 * self._inv_r0
 
-    @lazy_property
+    @property
     @implements(_galsim.spergel.Spergel.half_light_radius)
     def half_light_radius(self):
         return self._r0 * _spergel_hlr_pade(self.nu)
 
-    @lazy_property
+    @property
     def _shootxnorm(self):
         """Normalization for photon shooting"""
         return 1.0 / (2.0 * jnp.pi * jnp.power(2.0, self.nu) * _gammap1(self.nu))
 
-    @lazy_property
+    @property
     def _xnorm(self):
         """Normalization of xValue"""
         return self._shootxnorm * self.flux * self._inv_r0_sq
 
-    @lazy_property
+    @property
     def _xnorm0(self):
         """return z^nu K_nu(z) for z=0"""
         return jax.lax.select(
@@ -392,21 +391,21 @@ def __str__(self):
         s += ")"
         return s
 
-    @lazy_property
+    @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
 
-    @lazy_property
+    @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
 
-    @lazy_property
+    @property
     def _max_sb(self):
         # from SBSpergelImpl.h
         return jnp.abs(self._xnorm) * self._xnorm0
@@ -439,7 +438,7 @@ def withFlux(self, flux):
             gsparams=self.gsparams,
         )
 
-    @lazy_property
+    @property
     def _shoot_pos_cdf(self):
         zmax = calculateFluxRadius(
             1.0 - self.gsparams.shoot_accuracy, self.nu, zmax=30.0
@@ -459,7 +458,7 @@ def _shoot_pos(self, u):
         r = z * self._r0
         return r
 
-    @lazy_property
+    @property
     def _shoot_neg_cdf(self):
         # comment:
         # In the Galsim code the profile below rmin is linearized such that