fix: raise errors for celestial coords

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/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

tests/GalSim

@@ -98,6 +98,9 @@ def pytest_collection_modifyitems(config, items):
         ):
             item.add_marker(skip)
 
+        if any([t in item.nodeid for t in test_config["skipped_tests"]["coord"]]):
+            item.add_marker(skip)
+
 
 @lru_cache(maxsize=128)
 def _infile(val, fname):

tests/conftest.py

@@ -28,6 +28,13 @@ enabled_tests:
     - test_astropy.py
     - test_celestial.py
 
+# tests to explicitly skip
+# applied on top of the enabled set above
+skipped_tests:
+  coord:
+    - "tests/Coord/tests/test_celestial.py::test_xyz_raises"
+    - "tests/Coord/tests/test_celestial.py::test_greatcircle_raises"
+
 # This documents which error messages will be allowed
 # without being reported as an error. These typically
 # correspond to features that are not implemented yet
@@ -88,9 +95,8 @@ allowed_failures:
   - "module 'jax_galsim' has no attribute 'CorrelatedNoise'"
   - "CelestialCoord.precess is too slow"  # cannot get jax to warmup but once it does it passes
   - "ValueError not raised by from_xyz"
-  - "ValueError not raised by greatCirclePoint"
-  - "TypeError not raised by __mul__"
-  - "ValueError not raised by CelestialCoord"
+  # - "ValueError not raised by greatCirclePoint"
+  # - "ValueError not raised by CelestialCoord"
   - "module 'jax_galsim' has no attribute 'BaseCorrelatedNoise'"
   - "module 'jax_galsim' has no attribute 'fft'"
   - "Transform does not support callable arguments."

tests/galsim_tests_config.yaml

@@ -1,4 +1,5 @@
 import galsim as _galsim
+import jax.numpy as jnp
 import numpy as np
 import pytest
 
@@ -118,3 +119,22 @@ def test_celestial_jax_ecliptic_obliquity():
             ecliptic_obliquity(epoch).rad,
             _ecliptic_obliquity(epoch).rad,
         )
+
+
+def test_celestial_jax_xyz_raises():
+    np.testing.assert_raises(
+        Exception, jax_galsim.CelestialCoord.from_xyz, 0.0, 0.0, 0.0
+    )
+
+
+def test_celestial_jax_greatcircle_raises():
+    theta = 50 * jax_galsim.radians
+    eq1 = jax_galsim.CelestialCoord(
+        0 * jax_galsim.radians, 0 * jax_galsim.radians
+    )  # point on the equator
+    eq2 = jax_galsim.CelestialCoord(
+        jnp.array(1) * jax_galsim.radians, 0 * jax_galsim.radians
+    )  # 1 radian along equator
+
+    np.testing.assert_raises(Exception, eq1.greatCirclePoint, eq1, theta)
+    np.testing.assert_raises(Exception, eq2.greatCirclePoint, eq2, theta)