test: update tests with latest commits from main (#220)

* test: update tests to branch to ensure they pass

* fix: use in-place ops properly

* fix: buggy thing

* fix: correct FFT shape for inplace

* fix: avoid multiple extra calls to set things inplace

* test: fix more test warnings

* test: fix more test warnings

* test: latest commit

* fix: tests for implements maybe

* test: make tests robust across versions

* fix: ensure implements removes leading spaces in older pythons

* fix: make sure coeffs is cast to an array

* debug: stop on first failure for now

* test: use latest testing branch

* fix: ensure we cast to native byte order

* test: use latest testing code

* fix: more native byte order casts

* fix: make a copy

* style: pre-commit

* Apply suggestion from @beckermr

* test: try a new test that might be more robust

* test: run it all via split

* test: try this test

* test: use latest changes

* fix: do not store durations for float32 tests

* test: update tests submodule

Author: beckermr

.github/workflows/python_package.yaml

@@ -61,6 +61,17 @@ jobs:
             cat .test_durations*
           fi
 
+      - name: Test with pytest in float32
+        run: |
+          pytest \
+            -vv \
+            --durations=100 \
+            --randomly-seed=42 \
+            --splits ${NUM_SPLITS} --group ${{ matrix.group }} \
+            --splitting-algorithm least_duration \
+            --retries 1 \
+            --test-in-float32
+
       - name: Test with pytest
         run: |
           pytest \
@@ -74,13 +85,6 @@ jobs:
             --clean-durations \
             --retries 1
 
-      - name: Test with pytest in float32
-        if: ${{ matrix.group == '1' }}
-        run: |
-          pytest \
-            -vv \
-            --test-in-float32
-
       - name: Upload test durations
         uses: actions/upload-artifact@v7
         with:

jax_galsim/__init__.py

@@ -8,7 +8,7 @@
 from .errors import GalSimKeyError, GalSimIndexError, GalSimNotImplementedError
 from .errors import GalSimBoundsError, GalSimUndefinedBoundsError, GalSimImmutableError
 from .errors import GalSimIncompatibleValuesError, GalSimSEDError, GalSimHSMError
-from .errors import GalSimFFTSizeError
+from .errors import GalSimFFTSizeError, GalSimFFTSizeWarning
 from .errors import GalSimConfigError, GalSimConfigValueError
 from .errors import GalSimWarning, GalSimDeprecationWarning
 

jax_galsim/convolve.py

@@ -486,10 +486,12 @@ def _kValue(self, pos):
 
     def _drawKImage(self, image, jac=None):
         image = self.orig_obj._drawKImage(image, jac)
-        image._array = jnp.where(
-            jnp.abs(image.array) > self._min_acc_kvalue,
-            1.0 / image.array,
-            self._inv_min_acc_kvalue,
+        image._array = image._array.at[...].set(
+            jnp.where(
+                jnp.abs(image.array) > self._min_acc_kvalue,
+                1.0 / image.array,
+                self._inv_min_acc_kvalue,
+            )
         )
         kx, ky = image.get_pixel_centers()
         _jac = jnp.eye(2) if jac is None else jac
@@ -500,10 +502,12 @@ def _drawKImage(self, image, jac=None):
         )
         ksq = (kx**2 + ky**2) * image.scale**2
         # Set to zero outside of nominal maxk so as not to amplify high frequencies.
-        image._array = jnp.where(
-            ksq > self.maxk**2,
-            0.0,
-            image.array,
+        image._array = image._array.at[...].set(
+            jnp.where(
+                ksq > self.maxk**2,
+                0.0,
+                image.array,
+            )
         )
         return image
 

jax_galsim/core/utils.py

@@ -9,6 +9,17 @@
 from jax.tree_util import tree_flatten
 
 
+def cast_numpy_array_to_native_byte_order(arr):
+    """Cast an array to native byte order."""
+    if not isinstance(arr, np.ndarray):
+        return arr
+
+    if arr.dtype.isnative:
+        return arr
+
+    return arr.astype(arr.dtype.newbyteorder("="))
+
+
 def has_tracers(x):
     """Return True if the input item is a JAX tracer or object, False otherwise."""
     for item in tree_flatten(x)[0]:
@@ -296,7 +307,7 @@ class ParsedDoc(NamedTuple):
     sections: dict[str, str] = {}
 
 
-def _break_off_body_section_by_newline(body):
+def _break_off_body_section_by_newline(body, double_check_first_indent=False):
     first_lines = []
     body_lines = []
     found_first_break = False
@@ -314,7 +325,14 @@ def _break_off_body_section_by_newline(body):
         else:
             first_lines.append(line)
 
+    if double_check_first_indent and len(first_lines) > 1:
+        len_first_indent = len(first_lines[1]) - len(first_lines[1].lstrip())
+        if len_first_indent > 0:
+            first_indent = first_lines[1][:len_first_indent]
+            first_lines[0] = first_indent + first_lines[0].lstrip()
+
     firstline = "\n".join(first_lines)
+    firstline = textwrap.dedent(firstline)
     body = "\n".join(body_lines)
     body = textwrap.dedent(body.lstrip("\n"))
 
@@ -337,7 +355,9 @@ def _parse_galsimdoc(docstr):
 
     signature, body = "", docstr
 
-    firstline, body = _break_off_body_section_by_newline(body)
+    firstline, body = _break_off_body_section_by_newline(
+        body, double_check_first_indent=True
+    )
 
     summary = firstline
     if not summary:

jax_galsim/errors.py

@@ -5,6 +5,7 @@
     GalSimDeprecationWarning,
     GalSimError,
     GalSimFFTSizeError,
+    GalSimFFTSizeWarning,
     GalSimHSMError,
     GalSimImmutableError,
     GalSimIncompatibleValuesError,

jax_galsim/gsobject.py

@@ -809,7 +809,7 @@ def drawReal(self, image, add_to_image=False):
         im1 = self._drawReal(image)
         temp = im1.subImage(image.bounds)
         if add_to_image:
-            image._array = image._array + temp._array
+            image._array = image._array.at[...].add(temp._array)
         else:
             image._array = temp._array
 
@@ -929,7 +929,7 @@ def drawFFT_finish(self, image, kimage, wrap_size, add_to_image):
         # Add (a portion of) this to the original image.
         temp = real_image.subImage(image.bounds)
         if add_to_image:
-            image._array = image._array + temp._array
+            image._array = image._array.at[...].add(temp._array)
         else:
             image._array = temp._array
 
@@ -1043,7 +1043,7 @@ def drawKImage(
         if not add_to_image:
             image._array = im2._array
         else:
-            image._array = im2._array + image._array
+            image._array = image._array.at[...].add(im2._array)
 
         image_in._array = image._array
         image_in._bounds = image._bounds