ENH add photon shooting

jax_galsim/__init__.py

@@ -55,6 +55,7 @@
 from .sum import Add, Sum
 from .transform import Transform, Transformation
 from .convolve import Convolve, Convolution, Deconvolution, Deconvolve
+from .deltafunction import DeltaFunction
 
 # WCS
 from .wcs import (
@@ -68,6 +69,7 @@
 )
 from .fits import FitsHeader
 from .celestial import CelestialCoord
+from .fitswcs import TanWCS, FitsWCS, GSFitsWCS
 
 # Shear
 from .shear import Shear, _Shear
@@ -85,6 +87,10 @@
 )
 from .interpolatedimage import InterpolatedImage, _InterpolatedImage
 
+# Photon Shooting
+from .photon_array import PhotonArray
+from .sensor import Sensor
+
 # packages kept separate
 from . import bessel
 from . import fits

jax_galsim/angle.py

@@ -22,7 +22,7 @@
 from jax._src.numpy.util import _wraps
 from jax.tree_util import register_pytree_node_class
 
-from jax_galsim.core.utils import cast_scalar_to_float, ensure_hashable
+from jax_galsim.core.utils import cast_to_array_scalar, ensure_hashable
 
 
 @_wraps(_galsim.AngleUnit)
@@ -34,7 +34,7 @@ def __init__(self, value):
         """
         :param value:   The measure of the unit in radians.
         """
-        self._value = cast_scalar_to_float(value)
+        self._value = cast_to_array_scalar(value, dtype=float)
 
     @property
     def value(self):
@@ -142,7 +142,7 @@ def __init__(self, theta, unit=None):
             raise TypeError("Invalid unit %s of type %s" % (unit, type(unit)))
         else:
             # Normal case
-            self._rad = cast_scalar_to_float(theta) * unit.value
+            self._rad = cast_to_array_scalar(theta, dtype=float) * unit.value
 
     @property
     def rad(self):

jax_galsim/bounds.py

@@ -4,11 +4,7 @@
 from jax._src.numpy.util import _wraps
 from jax.tree_util import register_pytree_node_class
 
-from jax_galsim.core.utils import (
-    cast_scalar_to_float,
-    cast_scalar_to_int,
-    ensure_hashable,
-)
+from jax_galsim.core.utils import cast_to_float, cast_to_int, ensure_hashable
 from jax_galsim.position import Position, PositionD, PositionI
 
 
@@ -264,10 +260,10 @@ class BoundsD(Bounds):
 
     def __init__(self, *args, **kwargs):
         self._parse_args(*args, **kwargs)
-        self.xmin = cast_scalar_to_float(self.xmin)
-        self.xmax = cast_scalar_to_float(self.xmax)
-        self.ymin = cast_scalar_to_float(self.ymin)
-        self.ymax = cast_scalar_to_float(self.ymax)
+        self.xmin = cast_to_float(self.xmin)
+        self.xmax = cast_to_float(self.xmax)
+        self.ymin = cast_to_float(self.ymin)
+        self.ymax = cast_to_float(self.ymax)
 
     def _check_scalar(self, x, name):
         try:
@@ -298,10 +294,10 @@ class BoundsI(Bounds):
 
     def __init__(self, *args, **kwargs):
         self._parse_args(*args, **kwargs)
-        self.xmin = cast_scalar_to_int(self.xmin)
-        self.xmax = cast_scalar_to_int(self.xmax)
-        self.ymin = cast_scalar_to_int(self.ymin)
-        self.ymax = cast_scalar_to_int(self.ymax)
+        self.xmin = cast_to_int(self.xmin)
+        self.xmax = cast_to_int(self.xmax)
+        self.ymin = cast_to_int(self.ymin)
+        self.ymax = cast_to_int(self.ymax)
 
     def _check_scalar(self, x, name):
         try:

jax_galsim/box.py

@@ -6,6 +6,7 @@
 from jax_galsim.core.draw import draw_by_kValue, draw_by_xValue
 from jax_galsim.core.utils import ensure_hashable
 from jax_galsim.gsobject import GSObject
+from jax_galsim.random import UniformDeviate
 
 
 @_wraps(_galsim.Box)
@@ -115,6 +116,15 @@ def tree_unflatten(cls, aux_data, children):
             **aux_data
         )
 
+    @_wraps(_galsim.Box._shoot)
+    def _shoot(self, photons, rng):
+        ud = UniformDeviate(rng)
+
+        # this does not fill arrays like in galsim
+        photons.x = (ud.generate(photons.x) - 0.5) * self.width
+        photons.y = (ud.generate(photons.y) - 0.5) * self.height
+        photons.flux = self.flux / photons.size()
+
 
 @_wraps(_galsim.Pixel)
 @register_pytree_node_class

jax_galsim/convolve.py

@@ -6,6 +6,7 @@
 
 from jax_galsim.gsobject import GSObject
 from jax_galsim.gsparams import GSParams
+from jax_galsim.photon_array import PhotonArray
 
 
 @_wraps(
@@ -308,7 +309,15 @@ def _drawReal(self, image, jac=None, offset=(0.0, 0.0), flux_scaling=1.0):
         raise NotImplementedError("Real-space convolutions are not implemented")
 
     def _shoot(self, photons, rng):
-        raise NotImplementedError("Photon shooting convolutions are not implemented")
+        self.obj_list[0]._shoot(photons, rng)
+        # It may be necessary to shuffle when convolving because we do not have a
+        # gaurantee that the convolvee's photons are uncorrelated, e.g., they might
+        # both have their negative ones at the end.
+        # However, this decision is now made by the convolve method.
+        for obj in self.obj_list[1:]:
+            p1 = PhotonArray(len(photons))
+            obj._shoot(p1, rng)
+            photons.convolve(p1, rng)
 
     def _drawKImage(self, image, jac=None):
         image = self.obj_list[0]._drawKImage(image, jac)