ENH add rest of galsim test suite

jax_galsim/__init__.py

@@ -44,6 +44,7 @@
     ImageS,
     ImageUI,
     ImageUS,
+    _Image,
 )
 
 # GSObject
@@ -55,6 +56,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 (
@@ -86,6 +88,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_to_float_array_scalar, ensure_hashable
+from jax_galsim.core.utils import cast_to_float, ensure_hashable
 
 
 @_wraps(_galsim.AngleUnit)
@@ -34,7 +34,11 @@ def __init__(self, value):
         """
         :param value:   The measure of the unit in radians.
         """
-        self._value = cast_to_float_array_scalar(value)
+        if isinstance(value, AngleUnit):
+            raise TypeError("Cannot construct AngleUnit from another AngleUnit")
+        self._value = 1.0 * cast_to_float(
+            value
+        )  # this will cause an exception if things are not numeric
 
     @property
     def value(self):
@@ -142,7 +146,7 @@ def __init__(self, theta, unit=None):
             raise TypeError("Invalid unit %s of type %s" % (unit, type(unit)))
         else:
             # Normal case
-            self._rad = cast_to_float_array_scalar(theta) * unit.value
+            self._rad = cast_to_float(theta) * unit.value
 
     @property
     def rad(self):