TST fix tests for JAX-Galsim

.github/workflows/ci.yml

@@ -1,5 +1,9 @@
 name: GalSim CI
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
 on:
     push:
         branches:

setup.py

@@ -244,7 +244,7 @@ def supports_gpu(compiler, cc_type):
     extra_cflags = copt[cc_type]
     extra_lflags = lopt[cc_type]
     return try_compile(cpp_code, compiler, extra_cflags, extra_lflags)
-    
+
 # Check for the fftw3 library in some likely places
 def find_fftw_lib(output=False):
     import distutils.sysconfig
@@ -434,29 +434,29 @@ def find_eigen_dir(output=False):
         if output:
             print("Downloaded %s.  Unpacking tarball."%fname)
         with tarfile.open(fname) as tar:
-            
+
             def is_within_directory(directory, target):
-                
+
                 abs_directory = os.path.abspath(directory)
                 abs_target = os.path.abspath(target)
-            
+
                 prefix = os.path.commonprefix([abs_directory, abs_target])
-                
+
                 return prefix == abs_directory
-            
+
             def safe_extract(tar, path=".", members=None, *, numeric_owner=False):
                 # Avoid security vulnerability in tar.extractall function.
                 # This bit of code was added by the Advanced Research Center at Trellix in PR #1188.
                 # For more information about the security vulnerability, see
                 # https://github.com/advisories/GHSA-gw9q-c7gh-j9vm
-            
+
                 for member in tar.getmembers():
                     member_path = os.path.join(path, member.name)
                     if not is_within_directory(path, member_path):
                         raise Exception("Attempted Path Traversal in Tar File")
-            
-                tar.extractall(path, members, numeric_owner=numeric_owner) 
-            
+
+                tar.extractall(path, members, numeric_owner=numeric_owner)
+
             safe_extract(tar, dir)
         os.remove(fname)
         # This actually extracts into a subdirectory with a name eigen-eigen-5a0156e40feb/

tests/galsim_test_helpers.py

@@ -25,6 +25,31 @@
 from numpy.testing import assert_raises
 from numpy.testing import assert_warns
 
+__all__ = [
+    "default_params",
+    "gsobject_compare",
+    "printval",
+    "convertToShear",
+    "check_basic_x",
+    "check_basic_k",
+    "assert_floatlike",
+    "assert_intlike",
+    "check_basic",
+    "do_shoot",
+    "do_kvalue",
+    "radial_integrate",
+    "drawNoise",
+    "check_pickle",
+    "check_all_diff",
+    "timer",
+    "CaptureLog",
+    "assert_raises",
+    "assert_warns",
+    "Profile",
+    "galsim_backend",
+    "is_jax_galsim",
+]
+
 # This file has some helper functions that are used by tests from multiple files to help
 # avoid code duplication.
 
@@ -43,6 +68,18 @@
         integration_relerr = 1.e-6,
         integration_abserr = 1.e-8)
 
+
+def galsim_backend():
+    if "jax_galsim/__init__.py" in galsim.__file__:
+        return "jax_galsim"
+    else:
+        return "galsim"
+
+
+def is_jax_galsim():
+    return galsim_backend() == "jax_galsim"
+
+
 def gsobject_compare(obj1, obj2, conv=None, decimal=10):
     """Helper function to check that two GSObjects are equivalent
     """
@@ -133,20 +170,32 @@ def check_basic_x(prof, name, approx_maxsb=False, scale=None):
         np.testing.assert_allclose(
                 image(i,j), prof._xValue(galsim.PositionD(x,y)), rtol=1.e-5,
                 err_msg="%s profile sb image does not match _xValue at %d,%d"%(name,i,j))
-    assert prof.withFlux.__doc__ == galsim.GSObject.withFlux.__doc__
-    assert prof.__class__.withFlux.__doc__ == galsim.GSObject.withFlux.__doc__
+    if is_jax_galsim():
+        for line in galsim.GSObject.withFlux.__doc__.splitlines():
+            if line.strip() and "LAX" not in line:
+                assert line.strip() in prof.withFlux.__doc__, (
+                    prof.withFlux.__doc__, galsim.GSObject.withFlux.__doc__,
+                )
+        for line in galsim.GSObject.withFlux.__doc__.splitlines():
+            if line.strip() and "LAX" not in line:
+                assert line.strip() in prof.__class__.withFlux.__doc__, (
+                    prof.__class__.withFlux.__doc__, galsim.GSObject.withFlux.__doc__,
+                )
+    else:
+        assert prof.withFlux.__doc__ == galsim.GSObject.withFlux.__doc__
+        assert prof.__class__.withFlux.__doc__ == galsim.GSObject.withFlux.__doc__
 
     # Check negative flux:
     neg_image = prof.withFlux(-prof.flux).drawImage(method='sb', scale=scale, use_true_center=False)
-    np.testing.assert_almost_equal(neg_image.array/prof.flux, -image.array/prof.flux, 7,
-                                   '%s negative flux drawReal is not negative of +flux image'%name)
+    np.testing.assert_array_almost_equal(neg_image.array/prof.flux, -image.array/prof.flux, 7,
+                                         '%s negative flux drawReal is not negative of +flux image'%name)
 
     # Direct call to drawReal should also work and be equivalent to the above with scale = 1.
     prof.drawImage(image, method='sb', scale=1., use_true_center=False)
     image2 = image.copy()
     prof.drawReal(image2)
-    np.testing.assert_equal(image2.array, image.array,
-                            err_msg="%s drawReal not equivalent to drawImage"%name)
+    np.testing.assert_array_equal(image2.array, image.array,
+                                  err_msg="%s drawReal not equivalent to drawImage"%name)
 
     # If supposed to be axisymmetric, make sure it is.
     if prof.is_axisymmetric:
@@ -194,7 +243,7 @@ def check_basic_k(prof, name):
 
     # Check negative flux:
     neg_image = prof.withFlux(-prof.flux).drawKImage(kimage.copy())
-    np.testing.assert_almost_equal(neg_image.array/prof.flux, -kimage.array/prof.flux, 7,
+    np.testing.assert_array_almost_equal(neg_image.array/prof.flux, -kimage.array/prof.flux, 7,
                                    '%s negative flux drawK is not negative of +flux image'%name)
 
     # If supposed to be axisymmetric, make sure it is in the kValues.
@@ -206,6 +255,30 @@ def check_basic_k(prof, name):
             np.testing.assert_allclose(test_values, ref_value, rtol=1.e-5,
                                        err_msg="%s profile not axisymmetric in kValues"%name)
 
+def assert_floatlike(val):
+    assert (
+        isinstance(val, float)
+        or (
+            is_jax_galsim()
+            and hasattr(val, "shape")
+            and val.shape == ()
+            and hasattr(val, "dtype")
+            and val.dtype.name in ["float", "float32", "float64"]
+        )
+    ), "Value is not float-like: type(%r) = %r" % (val, type(val))
+
+def assert_intlike(val):
+    assert (
+        isinstance(val, int)
+        or (
+            is_jax_galsim()
+            and hasattr(val, "shape")
+            and val.shape == ()
+            and hasattr(val, "dtype")
+            and val.dtype.name in ["int", "int32", "int64"]
+        )
+    ), "Value is not int-like: type(%r) = %r" % (val, type(val))
+
 def check_basic(prof, name, approx_maxsb=False, scale=None, do_x=True, do_k=True):
     """Do some basic sanity checks that should work for all profiles.
     """
@@ -220,12 +293,12 @@ def check_basic(prof, name, approx_maxsb=False, scale=None, do_x=True, do_k=True
             prof.positive_flux - prof.negative_flux, prof.flux,
             err_msg="%s profile flux not equal to posflux + negflux"%name)
     assert isinstance(prof.centroid, galsim.PositionD)
-    assert isinstance(prof.flux, float)
-    assert isinstance(prof.positive_flux, float)
-    assert isinstance(prof.negative_flux, float)
-    assert isinstance(prof.max_sb, float)
-    assert isinstance(prof.stepk, float)
-    assert isinstance(prof.maxk, float)
+    assert_floatlike(prof.flux)
+    assert_floatlike(prof.positive_flux)
+    assert_floatlike(prof.negative_flux)
+    assert_floatlike(prof.max_sb)
+    assert_floatlike(prof.stepk)
+    assert_floatlike(prof.maxk)
     assert isinstance(prof.has_hard_edges, bool)
     assert isinstance(prof.is_axisymmetric, bool)
     assert isinstance(prof.is_analytic_x, bool)
@@ -298,6 +371,9 @@ def do_shoot(prof, img, name):
     print('nphot = ',nphot)
     img2 = img.copy()
 
+    if is_jax_galsim():
+        rtol *= 3
+
     # Use a deterministic random number generator so we don't fail tests because of rare flukes
     # in the random numbers.
     rng = galsim.UniformDeviate(12345)

tests/test_catalog.py

@@ -16,6 +16,7 @@
 #    and/or other materials provided with the distribution.
 #
 
+import os
 import numpy as np
 
 import galsim
@@ -198,7 +199,7 @@ def test_basic_dict():
 def test_single_row():
     """Test that we can read catalogs with just one row (#394)
     """
-    filename = "output/test394.txt"
+    filename = os.path.join(os.path.dirname(__file__), "output/test394.txt")
     with open(filename, 'w') as f:
         f.write("3 4 5\n")
     cat = galsim.Catalog(filename, file_type='ascii')

tests/test_celestial.py → tests/test_celestial_galsim.py

@@ -17,6 +17,7 @@
 #
 
 import numpy
+import numpy as np
 import os
 import math
 
@@ -122,11 +123,11 @@ def test_angle():
 
     # Check invalid constructors
     assert_raises(TypeError,galsim.AngleUnit, galsim.degrees)
-    assert_raises(ValueError,galsim.AngleUnit, 'spam')
+    assert_raises((ValueError, TypeError), galsim.AngleUnit, 'spam')
     assert_raises(TypeError,galsim.AngleUnit, 1, 3)
     assert_raises(TypeError,galsim.Angle, 3.4)
     assert_raises(TypeError,galsim.Angle, theta1, galsim.degrees)
-    assert_raises(ValueError,galsim.Angle, 'spam', galsim.degrees)
+    assert_raises((ValueError, TypeError), galsim.Angle, 'spam')
     assert_raises(TypeError,galsim.Angle, 1, 3)
 
 
@@ -155,7 +156,7 @@ def test_celestialcoord_basic():
 
     x, y, z = c1.get_xyz()
     print('c1 is at x,y,z = ',x,y,z)
-    np.testing.assert_equal((x,y,z), (1,0,0))
+    np.testing.assert_array_equal((x,y,z), (1,0,0))
     assert c1 == galsim.CelestialCoord.from_xyz(x,y,z)
 
     x, y, z = c2.get_xyz()
@@ -343,9 +344,19 @@ def test_projection():
 
     # First the trivial case
     p0 = center.project(center, projection='lambert')
-    assert p0 == (0.0 * galsim.arcsec, 0.0 * galsim.arcsec)
+    np.testing.assert_allclose(
+        (p0[0].rad, p0[1].rad),
+        (0.0, 0.0),
+        rtol=0,
+        atol=1e-16,
+    )
     c0 = center.deproject(*p0, projection='lambert')
-    assert c0 == center
+    np.testing.assert_allclose(
+        c0.rad,
+        center.rad,
+        rtol=0,
+        atol=1e-16,
+    )
     np.testing.assert_almost_equal(center.jac_deproject(*p0, projection='lambert').ravel(),
                                    (1,0,0,1))
 
@@ -398,9 +409,19 @@ def test_projection():
 
     # First the trivial case
     p0 = center.project(center, projection='stereographic')
-    assert p0 == (0.0 * galsim.arcsec, 0.0 * galsim.arcsec)
+    np.testing.assert_allclose(
+        (p0[0].rad, p0[1].rad),
+        (0.0, 0.0),
+        rtol=0,
+        atol=1e-16,
+    )
     c0 = center.deproject(*p0, projection='stereographic')
-    assert c0 == center
+    np.testing.assert_allclose(
+        c0.rad,
+        center.rad,
+        rtol=0,
+        atol=1e-16,
+    )
     np.testing.assert_almost_equal(center.jac_deproject(*p0, projection='stereographic').ravel(),
                                    (1,0,0,1))
 
@@ -456,9 +477,19 @@ def test_projection():
 
     # First the trivial case
     p0 = center.project(center, projection='gnomonic')
-    assert p0 == (0.0 * galsim.arcsec, 0.0 * galsim.arcsec)
+    np.testing.assert_allclose(
+        (p0[0].rad, p0[1].rad),
+        (0.0, 0.0),
+        rtol=0,
+        atol=1e-16,
+    )
     c0 = center.deproject(*p0, projection='gnomonic')
-    assert c0 == center
+    np.testing.assert_allclose(
+        c0.rad,
+        center.rad,
+        rtol=0,
+        atol=1e-16,
+    )
     np.testing.assert_almost_equal(center.jac_deproject(*p0, projection='gnomonic').ravel(),
                                    (1,0,0,1))
 
@@ -510,9 +541,19 @@ def test_projection():
 
     # First the trivial case
     p0 = center.project(center, projection='postel')
-    assert p0 == (0.0 * galsim.arcsec, 0.0 * galsim.arcsec)
+    np.testing.assert_allclose(
+        (p0[0].rad, p0[1].rad),
+        (0.0, 0.0),
+        rtol=0,
+        atol=1e-16,
+    )
     c0 = center.deproject(*p0, projection='postel')
-    assert c0 == center
+    np.testing.assert_allclose(
+        c0.rad,
+        center.rad,
+        rtol=0,
+        atol=1e-16,
+    )
     np.testing.assert_almost_equal(center.jac_deproject(*p0, projection='postel').ravel(),
                                    (1,0,0,1))
 

tests/test_chromatic.py

@@ -16,6 +16,7 @@
 #    and/or other materials provided with the distribution.
 #
 
+import copy
 import os
 import copy
 import numpy as np

tests/test_config_output.py

@@ -1238,7 +1238,7 @@ def test_config():
     }
 
     # Test yaml
-    yaml_file_name = "output/test_config.yaml"
+    yaml_file_name = os.path.join(os.path.dirname(__file__), "output/test_config.yaml")
     with open(yaml_file_name, 'w') as fout:
         yaml.dump(config, fout, default_flow_style=True)
     # String None will be coverted to a real None.  Convert here in the comparison dict
@@ -1252,7 +1252,7 @@ def test_config():
     assert config == dict(config2)
 
     # Test json
-    json_file_name = "output/test_config.json"
+    json_file_name = os.path.join(os.path.dirname(__file__), "output/test_config.json")
     with open(json_file_name, 'w') as fout:
         json.dump(config, fout)