fix: ensure bounds vmaps as best as possible, more equinox errors, types for __eq__, _BoundsD/I

.github/workflows/python_package.yaml

@@ -21,9 +21,9 @@ jobs:
       fail-fast: false
       matrix:
         python-version: ["3.12"]
-        group: [1, 2, 3, 4]
+        group: [1, 2, 3, 4, 5, 6]
     env:
-      NUM_SPLITS: 4
+      NUM_SPLITS: 6
 
     steps:
       - uses: actions/checkout@v6
@@ -57,24 +57,24 @@ jobs:
             cp .test_durations .test_durations.${{ matrix.group }}
             ls -lah .test_durations*
             echo " "
-            cat .test_durations*
           fi
 
       - name: Test with pytest in float32
         run: |
           pytest \
-            -vv \
+            -v \
             --durations=100 \
             --randomly-seed=42 \
             --splits ${NUM_SPLITS} --group ${{ matrix.group }} \
             --splitting-algorithm least_duration \
             --retries 1 \
-            --test-in-float32
+            --test-in-float32 \
+            -n 4
 
       - name: Test with pytest
         run: |
           pytest \
-            -vv \
+            -v \
             --durations=100 \
             --randomly-seed=42 \
             --splits ${NUM_SPLITS} --group ${{ matrix.group }} \
@@ -83,7 +83,7 @@ jobs:
             --splitting-algorithm least_duration \
             --clean-durations \
             --retries 1 \
-            -n 2
+            -n 4
 
       - name: Upload test durations
         uses: actions/upload-artifact@v7

docs/sharp-bits.rst

@@ -61,8 +61,21 @@ 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
---------------------------------------
+Fixed Array Shapes in JAX Function Transformations
+--------------------------------------------------
+
+JAX function transformations (e.g., ``jax.jit``, ``jax.vmap``, etc.) require statically known
+array shapes in order to support tracing. To support this, the JAX-GalSim ``BoundsI`` class must
+have a statically known shape. Further this class can be instantiated via the syntax
+``BoundsI(xmin=..., deltax=..., ymin=..., deltay=...)`` where ``deltax/y`` are the statically defined
+shape. ``BoundsI`` classes may have dynamically set ``x/ymin`` values. However, in this case the ``&``
+and ``+`` operations, which can change the shape of the ``BoundsI`` instance are not allowed in
+JAX-traced code. ``BoundsI`` instances have a special method ``isStatic()`` which returns ``True``
+if the object was instantiated with statically know ``x/ymin`` values. A static ``BoundsI`` class
+cannot be converted to a dynamic one via assignment and an attempt to do so will raise an exception.
+
+Scalar Types, Array Types, and Type Casting
+-------------------------------------------
 
 With the use of JAX, there are now many possible types for numeric data. These include
 
@@ -89,6 +102,24 @@ These rules allow JAX-GalSim to transparently handle JAX's tracing operations, b
 the code raising generic ``Exception`` instances instead of more specific ``GalSim`` exceptions in
 some cases.
 
+Object Comparison with the ``==`` Operator
+------------------------------------------
+
+In JAX-GalSim, all objects which define arrays to be traced by JAX will return JAX boolean
+array scalars (i.e., ``jax.numpy.array(True)`` or ``jax.numpy.array(False)``) as the result
+of the ``==`` operator. Otherwise the return value is a Python boolean. Important cases of this
+rule are static ``BoundsI`` objects, ``Interpolant`` objects (and their subclasses), and ``GSParams``
+objects, all of which return Python boolean values (i.e. ``True`` and ``False``). These difference
+can be a source of subtle bugs since the negation of JAX array boolean values is typically done
+with ``~``, while for Python boolean values it is done with ``not``. Mixing these two forms can
+cause unexpected and incorrect results since
+
+.. code-block:: python
+
+   >>> ~True is False
+   <python-input-0>:1: SyntaxWarning: "is" with 'int' literal. Did you mean "=="?
+   False
+
 Random Number Generation
 ------------------------
 

jax_galsim/__init__.py

@@ -33,7 +33,7 @@
 )
 
 # Basic building blocks
-from .bounds import Bounds, BoundsD, BoundsI
+from .bounds import Bounds, BoundsD, BoundsI, _BoundsD, _BoundsI
 from .gsparams import GSParams
 from .position import Position, PositionD, PositionI
 from .angle import Angle, AngleUnit, _Angle, radians, hours, degrees, arcmin, arcsec

jax_galsim/angle.py

@@ -112,10 +112,13 @@ def __repr__(self):
             return "galsim.AngleUnit(%r)" % (ensure_hashable(self.value),)
 
     def __eq__(self, other):
-        return isinstance(other, AngleUnit) and jnp.array_equal(self.value, other.value)
+        if not isinstance(other, AngleUnit):
+            return jnp.array(False)
+        else:
+            return jnp.array_equal(self.value, other.value)
 
     def __ne__(self, other):
-        return not self.__eq__(other)
+        return ~self.__eq__(other)
 
     def __hash__(self):
         return hash(("galsim.AngleUnit", ensure_hashable(self.value)))
@@ -253,10 +256,13 @@ def __repr__(self):
         return "galsim.Angle(%r, galsim.radians)" % (ensure_hashable(self.rad),)
 
     def __eq__(self, other):
-        return isinstance(other, Angle) and jnp.array_equal(self.rad, other.rad)
+        if not isinstance(other, Angle):
+            return jnp.array(False)
+        else:
+            return jnp.array_equal(self.rad, other.rad)
 
     def __ne__(self, other):
-        return not self.__eq__(other)
+        return ~self.__eq__(other)
 
     def __le__(self, other):
         if not isinstance(other, Angle):