Merge pull request #243 from GalSim-developers/equinox-err-2

Author: beckermr

jax_galsim/angle.py

@@ -21,9 +21,30 @@
 # SOFTWARE.
 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.utils import cast_to_float, ensure_hashable, implements
+from jax_galsim.core.utils import (
+    cast_to_float,
+    ensure_hashable,
+    has_tracers,
+    implements,
+)
+
+NON_COMPLEX_TYPES = (
+    float,
+    int,
+    np.int16,
+    np.int32,
+    np.int64,
+    np.float32,
+    np.float64,
+    jnp.int16,
+    jnp.int32,
+    jnp.int64,
+    jnp.float32,
+    jnp.float64,
+)
 
 
 @implements(_galsim.AngleUnit)
@@ -178,15 +199,23 @@ 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)):
+            raise TypeError(
+                "Cannot multiply Angle by %s of type %s" % (other, type(other))
+            )
         return _Angle(self._rad * other)
 
     __rmul__ = __mul__
 
     def __div__(self, other):
         if isinstance(other, AngleUnit):
             return self._rad / other.value
-        else:
+        elif has_tracers(other) or isinstance(other, NON_COMPLEX_TYPES):
             return _Angle(self._rad / other)
+        else:
+            raise TypeError(
+                "Cannot divide Angle by %s of type %s" % (other, type(other))
+            )
 
     __truediv__ = __div__
 

jax_galsim/bounds.py

@@ -1,34 +1,24 @@
 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
 
-CONST_TYPES = (float, int, np.ndarray, np.int32, np.int64, np.float32, np.float64)
-CONST_TYPES_WITH_JAX = CONST_TYPES + (
-    jax.Array,
-    jnp.array,
-    jnp.int32,
-    jnp.int64,
-    jnp.float32,
-    jnp.float64,
-)
-
-# TODO: write extra docs for JAX changes
 BOUNDS_LAX_DESCR = """\
 The JAX implementation
 
 - will not always test whether the bounds are valid
-- will not always test whether BoundsI is initialized with integers
 
 Further, the JAX implementation adds a new method, ``isStatic`` to the
 ``BoundsI`` class. If JAX-GalSim detects that the ``BoundsI`` instance
@@ -525,31 +515,27 @@ def __init__(self, *args, **kwargs):
                 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)
+        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))
 
-        if (self.deltax != int(self.deltax)) or (self.deltay != int(self.deltay)):
-            raise TypeError("BoundsI must be initialized with integer values")
+        if has_tracers(self._xmin) or has_tracers(self._ymin):
+            self._isstatic = False
+
+        # validate inputs are ints
+        self._xmin = check_is_int_then_cast(
+            self._xmin, "BoundsI must be initialized with integer values"
+        )
+        self._ymin = check_is_int_then_cast(
+            self._ymin, "BoundsI must be initialized with integer values"
+        )
 
         if self.deltax < 1 and self.deltay < 1:
             self._isdefined = False
 
-        # for simple inputs, we can check if the bounds are valid ints
-        if isinstance(self._xmin, CONST_TYPES) and self._xmin != int(self._xmin):
-            raise TypeError("BoundsI must be initialized with integer values")
-
-        if isinstance(self._ymin, CONST_TYPES) and self._ymin != int(self._ymin):
-            raise TypeError("BoundsI must be initialized with integer values")
-
-        if not has_tracers(self._xmin) and not has_tracers(self._ymin):
-            self._isstatic = True
-            self._xmin = int(np.trunc(self._xmin))
-            self._ymin = int(np.trunc(self._ymin))
-        else:
-            self._isstatic = False
-            self._xmin = cast_to_float(jnp.trunc(self._xmin))
-            self._ymin = cast_to_float(jnp.trunc(self._ymin))
-
         if force_static and not self._isstatic:
             raise RuntimeError(
                 "BoundsI initialized with non-static "

jax_galsim/celestial.py

@@ -23,9 +23,11 @@
 from functools import partial
 
 import coord as _coord
+import equinox
 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.angle import Angle, _Angle, arcsec, degrees, radians
@@ -74,6 +76,16 @@ def __init__(self, ra, dec=None):
         elif not isinstance(dec, Angle):
             raise TypeError("dec must be a galsim.Angle")
         else:
+            if isinstance(dec._rad, (float, int)):
+                if dec._rad < -np.pi / 2 or dec._rad > np.pi / 2:
+                    raise ValueError("dec must be between -90 deg and +90 deg.")
+            else:
+                dec._rad = equinox.error_if(
+                    jnp.array(dec._rad),
+                    jnp.any((dec._rad < -jnp.pi / 2) | (dec._rad > jnp.pi / 2)),
+                    "dec must be between -90 deg and +90 deg.",
+                )
+
             # Normal case
             self._ra = ra
             self._dec = dec
@@ -121,15 +133,14 @@ def get_xyz(self):
 
     @staticmethod
     @jax.jit
-    @implements(
-        _galsim.celestial.CelestialCoord.from_xyz,
-        lax_description=(
-            "The JAX version of this static method does not check that the norm of the input "
-            "vector is non-zero."
-        ),
-    )
+    @implements(_galsim.celestial.CelestialCoord.from_xyz)
     def from_xyz(x, y, z):
         norm = jnp.sqrt(x * x + y * y + z * z)
+        norm = equinox.error_if(
+            norm,
+            jnp.any(norm == 0),
+            "CelestialCoord for position (0,0,0) is undefined.",
+        )
         ret = CelestialCoord.__new__(CelestialCoord)
         ret._x = x / norm
         ret._y = y / norm
@@ -236,13 +247,7 @@ def distanceTo(self, coord2):
 
         return _Angle(theta)
 
-    @implements(
-        _galsim.celestial.CelestialCoord.greatCirclePoint,
-        lax_description=(
-            "The JAX version of this method does not check that coord2 defines a unique great "
-            "circle with the current coord at angle theta."
-        ),
-    )
+    @implements(_galsim.celestial.CelestialCoord.greatCirclePoint)
     @jax.jit
     def greatCirclePoint(self, coord2, theta):
         aux = self._get_aux()
@@ -280,8 +285,11 @@ def greatCirclePoint(self, coord2, theta):
 
         # Normalize
         wr = (wx**2 + wy**2 + wz**2) ** 0.5
-        # if wr == 0.:
-        #     raise ValueError("coord2 does not define a unique great circle with self.")
+        wr = equinox.error_if(
+            wr,
+            jnp.any(wr == 0),
+            "coord2 does not define a unique great circle with self.",
+        )
         wx /= wr
         wy /= wr
         wz /= wr

jax_galsim/core/utils.py

@@ -3,11 +3,60 @@
 from functools import partial
 from typing import NamedTuple
 
+import equinox
 import jax
 import jax.numpy as jnp
 import numpy as np
 from jax.tree_util import tree_flatten
 
+CONST_TYPES = (
+    float,
+    int,
+    np.ndarray,
+    np.int8,
+    np.int16,
+    np.int32,
+    np.int64,
+    np.float16,
+    np.float32,
+    np.float64,
+    np.complex64,
+    np.complex128,
+)
+CONST_TYPES_WITH_JAX = CONST_TYPES + (
+    jax.Array,
+    jnp.ndarray,
+    jnp.int8,
+    jnp.int16,
+    jnp.int32,
+    jnp.int64,
+    jnp.float32,
+    jnp.float64,
+    jnp.complex64,
+    jnp.complex128,
+)
+
+
+def check_is_int_then_cast(val, msg):
+    """Check if `val` is an integer, raise if not, otherwise cast to int."""
+    # for simple inputs, we can check direct in python
+    if isinstance(val, CONST_TYPES) and not has_tracers(val):
+        val = cast_to_python_float(val)
+        if val != int(val):
+            raise TypeError(msg)
+        val = int(val)
+    else:
+        # otherwise we use more opaque checking upon jit via equinox
+        val = jnp.array(val)
+        val = equinox.error_if(
+            val,
+            np.any(val != jnp.trunc(val)),
+            msg,
+        )
+        val = val.astype(int)
+
+    return val
+
 
 def cast_numpy_array_to_native_byte_order(arr):
     """Cast an array to native byte order."""

jax_galsim/fitswcs.py

@@ -1,6 +1,7 @@
 import copy
 import os
 
+import equinox
 import galsim as _galsim
 import jax
 import jax.numpy as jnp
@@ -1094,12 +1095,10 @@ def _step(i, args):
         unroll=True,
     )[0:4]
 
-    x, y = jax.lax.cond(
-        jnp.maximum(jnp.max(jnp.abs(dx)), jnp.max(jnp.abs(dy))) > 2e-12,
-        lambda x, y: (x * jnp.nan, y * jnp.nan),
-        lambda x, y: (x, y),
-        x,
-        y,
+    x, y = equinox.error_if(
+        (x, y),
+        jnp.any(jnp.maximum(jnp.max(jnp.abs(dx)), jnp.max(jnp.abs(dy))) > 2e-12),
+        "Unable to solve for image_pos (max iter reached).",
     )
 
     return x, y