test: clean up tests for draw

GalSim-developers · Aug 23, 2024 · 63855f9 · 63855f9
1 parent 682f592
commit 63855f9
Showing 1 changed file with 10 additions and 8 deletions.
diff --git a/tests/test_draw.py b/tests/test_draw.py
@@ -1259,10 +1259,11 @@ def test_fft():
                              [4,6,8,4],
                              [2,4,6,6] ],
                            xmin=-2, ymin=-2, dtype=dt, scale=0.1)
-        if is_jax_galsim() and dt not in [np.complex128, complex]:
-            kim = xim.calculate_fft()
-            xim2 = kim.calculate_inverse_fft()
-            np.testing.assert_array_almost_equal(xim.array, xim2.array)
+        if is_jax_galsim():
+            if dt not in [np.complex128, complex]:
+                kim = xim.calculate_fft()
+                xim2 = kim.calculate_inverse_fft()
+                np.testing.assert_array_almost_equal(xim.array, xim2.array)
         else:
             kim = xim.calculate_fft()
             xim2 = kim.calculate_inverse_fft()
@@ -1298,10 +1299,11 @@ def test_fft():
         xim2 = galsim.Image([ [2,4,6],
                               [4,6,8] ],
                             xmin=-2, ymin=-1, dtype=dt, scale=0.1)
-        if is_jax_galsim() and dt not in [np.complex128, complex]:
-            kim = xim.calculate_fft()
-            kim2 = xim2.calculate_fft()
-            np.testing.assert_array_almost_equal(kim.array, kim2.array)
+        if is_jax_galsim():
+            if dt not in [np.complex128, complex]:
+                kim = xim.calculate_fft()
+                kim2 = xim2.calculate_fft()
+                np.testing.assert_array_almost_equal(kim.array, kim2.array)
         else:
             kim = xim.calculate_fft()
             kim2 = xim2.calculate_fft()