Merge pull request #248 from GalSim-developers/typing-inits-cleanup

fix: clean up type handling

Author: beckermr

docs/sharp-bits.rst

@@ -61,6 +61,34 @@ does not affect the original.
    # JAX-GalSim — real_part is a copy
    real_part = complex_image.real  # independent array
 
+Scalar Types, Array Types, and Casting
+--------------------------------------
+
+With the use of JAX, there are now many possible types for numeric data. These include
+
+- **Python scalars**: Objects with types that are ``float``, ``int``, or ``complex``.
+- **NumPy scalars**: Objects with types that are subclasses of ``np.floating``, ``np.integer``, etc.
+- **NumPy array scalars**: Objects with a type that is ``np.ndarray`` and have ``np.ndim(...) == 0``.
+- **NumPy arrays**: Objects with a type that is ``np.ndarray`` and have ``np.ndim(...) > 0``.
+- **JAX array scalars**: Objects with a type that is ``jax.numpy.ndarray`` and have ``jax.numpy.ndim(...) == 0``.
+- **JAX arrays**: Objects with a type that is ``jax.numpy.ndarray`` and have ``jax.numpy.ndim(...) > 0``.
+
+**JAX does not have pure scalar types like NumPy. JAX uses array scalars for those instead.**
+
+JAX-GalSim uses the following rules when handling data types and casting.
+
+- If the item is a Python numeric type (i.e., ``int`` or ``float``) or a
+  NumPy scalar type (i.e., ``isinstance(x, np.number)``, ``isinstance(x, np.integer)``, etc.),
+  convert it to a Python type of the appropriate kind.
+- For all other array-like types, cast to the correct type via ``jax.numpy.astype(x, ...)``.
+- For putting data into FITS headers only, JAX-GalSim converts of NumPy/JAX arrays to Python
+  numeric types as long as there is one element in the array (i.e., it is a NumPy scalar type,
+  an array scalar, or a 1D array with one element).
+
+These rules allow JAX-GalSim to transparently handle JAX's tracing operations, but can result in
+the code raising generic ``Exception`` instances instead of more specific ``GalSim`` exceptions in
+some cases.
+
 Random Number Generation
 ------------------------
 
@@ -163,9 +191,6 @@ profile parameters passed into a ``jit``-compiled function):
    def good(sigma):
        return jax.lax.cond(sigma > 1.0, lambda s: s * 2, lambda s: s, sigma)
 
-JAX-GalSim uses an internal ``has_tracers()`` utility to detect tracing and
-avoid problematic control flow in its own implementations.
-
 Fixed output shapes
 ^^^^^^^^^^^^^^^^^^^
 
@@ -197,20 +222,9 @@ The ``__init__`` gotcha
 
 During ``jit`` tracing, JAX calls constructors with **tracer objects** rather
 than concrete Python numbers. Type checks like ``isinstance(sigma, float)`` will
-fail on tracers. JAX-GalSim handles this internally, but if you subclass any
-JAX-GalSim object, be aware that ``__init__`` may receive tracers:
-
-.. code-block:: python
-
-   from jax_galsim.core.utils import has_tracers
-
-   class MyProfile(jax_galsim.GSObject):
-       def __init__(self, sigma, gsparams=None):
-           if not has_tracers(sigma):
-               # Only validate with concrete values
-               if sigma <= 0:
-                   raise ValueError("sigma must be positive")
-           ...
+return ``False`` on tracers, and you cannot check correctness of values (e.g.,
+``if sigma > 0: ...```). JAX-GalSim handles this internally, but if you subclass any
+JAX-GalSim object, be aware that ``__init__`` may receive tracers.
 
 Profile Restrictions
 --------------------
@@ -221,12 +235,9 @@ Some GalSim features are not yet implemented in JAX-GalSim:
 - **ChromaticObject**: All chromatic functionality (wavelength-dependent
   profiles) is not available.
 - **InterpolatedKImage**: Not implemented.
-- **Airy, Kolmogorov, OpticalPSF, RealGalaxy**: See :doc:`api-coverage` for
+- **Airy, Kolmogorov, OpticalPSF, RealGalaxy, etc.**: See :doc:`api-coverage` for
   the full list.
 
-The project currently implements **22.5 %** of the GalSim public API, focused
-on the most commonly used profiles and operations.
-
 Numerical Precision
 -------------------
 
@@ -249,11 +260,11 @@ These differences are typically at the level of floating-point round-off
 should not affect scientific conclusions.
 
 ⚠️ Additional Sharp Bits
---------------------------
+------------------------
 
 In the :doc:`api/index` you will find **🔪 JAX-GalSim - The Sharp Bits 🔪** blocks highlighting additional important caveats for specific classes and or methods. These could include things like:
 
-- Many classes do not perform some of Galsim's test for correctness during initialization (e.g., :meth:`~jax_galsim.GSObject.drawImage`).
+- Some classes do not perform some of Galsim's test for correctness during initialization (e.g., :meth:`~jax_galsim.InterpolatedImage`).
 - Certain profiles might not be auto-differentiable with respect to some of their parameters (e.g., :class:`~jax_galsim.Spergel`, :class:`~jax_galsim.Moffat`)
 - Limitations regarding what types of inputes are handled (e.g., :meth:`~jax_galsim.Image.calculate_fft` does not accept complex dtypes.)
 

jax_galsim/angle.py

@@ -27,7 +27,6 @@
 from jax_galsim.core.utils import (
     cast_to_float,
     ensure_hashable,
-    has_tracers,
     implements,
 )
 
@@ -55,7 +54,7 @@ class AngleUnit(object):
     def __init__(self, value):
         if isinstance(value, AngleUnit):
             raise TypeError("Cannot construct AngleUnit from another AngleUnit")
-        self._value = cast_to_float(value)
+        self._value = cast_to_float(value, accept_strings=True)
 
     @property
     @implements(_galsim.AngleUnit.value)
@@ -199,7 +198,7 @@ def __sub__(self, other):
         return _Angle(self._rad - other._rad)
 
     def __mul__(self, other):
-        if not (has_tracers(other) or isinstance(other, NON_COMPLEX_TYPES)):
+        if isinstance(other, (Angle, AngleUnit)):
             raise TypeError(
                 "Cannot multiply Angle by %s of type %s" % (other, type(other))
             )
@@ -210,7 +209,7 @@ def __mul__(self, other):
     def __div__(self, other):
         if isinstance(other, AngleUnit):
             return self._rad / other.value
-        elif has_tracers(other) or isinstance(other, NON_COMPLEX_TYPES):
+        elif not isinstance(other, Angle):
             return _Angle(self._rad / other)
         else:
             raise TypeError(

jax_galsim/bounds.py

@@ -1,16 +1,14 @@
 import galsim as _galsim
 import jax
 import jax.numpy as jnp
+import numpy as np
 from jax.tree_util import register_pytree_node_class
 
 from jax_galsim.core.utils import (
-    CONST_TYPES,
     cast_to_float,
     cast_to_int,
-    cast_to_python_float,
     check_is_int_then_cast,
     ensure_hashable,
-    has_tracers,
     implements,
 )
 from jax_galsim.position import Position, PositionD, PositionI
@@ -23,10 +21,7 @@
 Further, the JAX implementation adds a new method, ``isStatic`` to the
 ``BoundsI`` class. If JAX-GalSim detects that the ``BoundsI`` instance
 has been instantiated with static, known values, ``isStatic()`` will
-return ``True``. You can indicate to JAX-GalSim that a ``BoundsI``
-instance should be static via initializing it with the ``static``
-keyword set to the ``True``. If the object detects that it is being
-initialized with non-static data, an error will be raised.
+return ``True``.
 
 ``BoundsI`` objects in JAX-Galsim support an additional initialization
 call ``BoundsI(xmin=..., deltax=..., ymin=..., deltay=...)``. In this case,
@@ -139,8 +134,8 @@ def _parse_args(self, *args, **kwargs):
         else:
             max_delta = 1
         if (
-            isinstance(self.deltax, CONST_TYPES)
-            and isinstance(self.deltay, CONST_TYPES)
+            isinstance(self.deltax, (int, float, np.integer, np.floating))
+            and isinstance(self.deltay, (int, float, np.integer, np.floating))
             and (self.deltax < max_delta or self.deltay < max_delta)
         ):
             self._isdefined = False
@@ -364,10 +359,8 @@ def from_galsim(cls, galsim_bounds):
         """Create a jax_galsim `BoundsD/I` from a `galsim.BoundsD/I` object."""
         if isinstance(galsim_bounds, _galsim.BoundsD):
             _cls = BoundsD
-            kwargs = {}
         elif isinstance(galsim_bounds, _galsim.BoundsI):
             _cls = BoundsI
-            kwargs = {"static": True}
         else:
             raise TypeError(
                 "galsim_bounds must be either a %s or a %s"
@@ -379,7 +372,6 @@ def from_galsim(cls, galsim_bounds):
                 galsim_bounds.xmax,
                 galsim_bounds.ymin,
                 galsim_bounds.ymax,
-                **kwargs,
             )
         else:
             return _cls()
@@ -505,24 +497,25 @@ def __init__(self, *args, **kwargs):
         # initial setting to let stuff pass through freely
         self._isstatic = True
 
-        force_static = kwargs.pop("static", False)
-
         self._parse_args(*args, **kwargs)
 
-        if has_tracers(self.deltax) or has_tracers(self.deltay):
-            raise RuntimeError(
-                "Jax-GalSim BoundsI instances must have a fixed width! "
-                f"Got deltax,deltay = {self.deltax!r},{self.deltay!r}."
-            )
-
-        self.deltax = cast_to_python_float(self.deltax)
-        self.deltay = cast_to_python_float(self.deltay)
+        self.deltax = cast_to_float(self.deltax)
+        self.deltay = cast_to_float(self.deltay)
         if (self.deltax != int(self.deltax)) or (self.deltay != int(self.deltay)):
             raise TypeError("BoundsI must be initialized with integer values")
-        self.deltax = int(cast_to_int(self.deltax))
-        self.deltay = int(cast_to_int(self.deltay))
+        self.deltax = cast_to_int(self.deltax)
+        self.deltay = cast_to_int(self.deltay)
 
-        if has_tracers(self._xmin) or has_tracers(self._ymin):
+        if not (
+            isinstance(
+                self._xmin,
+                (int, float, np.floating, np.integer),
+            )
+            and isinstance(
+                self._ymin,
+                (int, float, np.floating, np.integer),
+            )
+        ):
             self._isstatic = False
 
         # validate inputs are ints
@@ -536,13 +529,6 @@ def __init__(self, *args, **kwargs):
         if self.deltax < 1 and self.deltay < 1:
             self._isdefined = False
 
-        if force_static and not self._isstatic:
-            raise RuntimeError(
-                "BoundsI initialized with non-static "
-                f"data (xmin,ymin = {self._xmin},{self._yminb}) "
-                "when static data was explicitly requested."
-            )
-
     def _check_scalar(self, x, name):
         try:
             if (