gh-1076: Use array comparison functions from array-api-extra

pyproject.toml

@@ -14,6 +14,7 @@ dev = [
     "coverage[toml]",
     "hatch-vcs",
     "hatchling",
+    "jupyter",
     "nox",
     "nox-uv",
     {include-group = "build"},

tests/benchmarks/test_arraytools.py

@@ -4,15 +4,15 @@
 
 import pytest
 
+import array_api_extra as xpx
+
 import glass.arraytools
 
 if TYPE_CHECKING:
     from types import ModuleType
 
     from pytest_benchmark.fixture import BenchmarkFixture
 
-    from tests.fixtures.helper_classes import Compare
-
 
 @pytest.mark.unstable
 def test_broadcast_leading_axes(
@@ -44,7 +44,6 @@ def test_broadcast_leading_axes(
 @pytest.mark.unstable
 def test_cumulative_trapezoid_1d(
     benchmark: BenchmarkFixture,
-    compare: type[Compare],
     xpb: ModuleType,
 ) -> None:
     """Benchmark test for glass.arraytools.cumulative_trapezoid."""
@@ -56,7 +55,7 @@ def test_cumulative_trapezoid_1d(
     ct = benchmark(glass.arraytools.cumulative_trapezoid, f, x)
 
     # Compare to int64 as old versions of glass round to int64 if `dtype` is not passed.
-    compare.assert_allclose(
+    xpx.testing.assert_close(
         xpb.asarray(ct[:4], dtype=xpb.int64),
         xpb.asarray([0, 1, 4, 7]),
     )
@@ -66,7 +65,6 @@ def test_cumulative_trapezoid_1d(
 @pytest.mark.unstable
 def test_cumulative_trapezoid_2d(
     benchmark: BenchmarkFixture,
-    compare: type[Compare],
     xpb: ModuleType,
 ) -> None:
     """Benchmark test for glass.arraytools.cumulative_trapezoid."""
@@ -80,16 +78,16 @@ def test_cumulative_trapezoid_2d(
     )
     x = xpb.arange(scaled_length)  # [0, 1, 2, 3,...]
 
-    expected_first_4_out = [0, 1, 4, 7]
-
     ct = benchmark(glass.arraytools.cumulative_trapezoid, f, x)
 
+    expected_first_4_out = xpb.asarray([0, 1, 4, 7])
+
     # Compare to int64 as old versions of glass round to int64 if `dtype` is not passed.
-    compare.assert_allclose(
+    xpx.testing.assert_close(
         xpb.asarray(ct[0, :4], dtype=xpb.int64),
         expected_first_4_out,
     )
-    compare.assert_allclose(
+    xpx.testing.assert_close(
         xpb.asarray(ct[1, :4], dtype=xpb.int64),
         expected_first_4_out,
     )

tests/benchmarks/test_fields.py

@@ -5,6 +5,8 @@
 
 import pytest
 
+import array_api_extra as xpx
+
 import glass
 import glass.fields
 import glass.healpix as hp
@@ -16,7 +18,7 @@
     from pytest_benchmark.fixture import BenchmarkFixture
 
     from glass._types import AngularPowerSpectra, UnifiedGenerator
-    from tests.fixtures.helper_classes import Compare, GeneratorConsumer
+    from tests.fixtures.helper_classes import GeneratorConsumer
 
 
 @pytest.mark.stable
@@ -102,7 +104,6 @@ def function_to_benchmark() -> list[Any]:
 @pytest.mark.stable
 def test_iternorm_k_0(
     benchmark: BenchmarkFixture,
-    compare: type[Compare],
     generator_consumer: type[GeneratorConsumer],
     xpb: ModuleType,
 ) -> None:
@@ -119,13 +120,12 @@ def function_to_benchmark() -> list[Any]:
     j, a, s = results[0]
     assert j is None
     assert a.shape == (0,)
-    compare.assert_allclose(xpb.asarray(s), 1.0)
+    xpx.testing.assert_close(s, xpb.asarray(1.0))
 
 
 @pytest.mark.stable
 def test_cls2cov(
     benchmark: BenchmarkFixture,
-    compare: type[Compare],
     generator_consumer: type[GeneratorConsumer],
     xpb: ModuleType,
 ) -> None:
@@ -148,9 +148,9 @@ def function_to_benchmark() -> list[Any]:
     assert cov.shape == (nl, nc + 1)
     assert cov.dtype == xpb.float64
 
-    compare.assert_allclose(cov[:, 0], xpb.asarray([1.0, 1.5, 2.0]))
-    compare.assert_allclose(cov[:, 1], xpb.asarray([1.5, 2.0, 2.5]))
-    compare.assert_allclose(cov[:, 2], 0)
+    xpx.testing.assert_close(cov[:, 0], xpb.asarray([1.0, 1.5, 2.0]))
+    xpx.testing.assert_close(cov[:, 1], xpb.asarray([1.5, 2.0, 2.5]))
+    xpx.testing.assert_close(cov[:, 2], xpb.asarray(0.0), check_shape=False)
 
 
 @pytest.mark.stable
@@ -185,7 +185,6 @@ def function_to_benchmark() -> list[Any]:
 @pytest.mark.parametrize("ncorr", [None, 1])
 def test_generate(
     benchmark: BenchmarkFixture,
-    compare: type[Compare],
     generator_consumer: type[GeneratorConsumer],
     xpb: ModuleType,
     ncorr: int | None,
@@ -214,13 +213,12 @@ def function_to_benchmark() -> list[Any]:
 
     for field in result:
         assert field.shape == (hp.nside2npix(nside),)
-    compare.assert_allclose(result[1], result[0] ** 2, atol=1e-05)
+    xpx.testing.assert_close(result[1], result[0] ** 2, atol=1e-05)
 
 
 @pytest.mark.unstable
 def test_getcl_lmax_0(
     benchmark: BenchmarkFixture,
-    compare: type[Compare],
     xpb: ModuleType,
 ) -> None:
     """Benchmarks for glass.getcl with lmax of 0."""
@@ -245,13 +243,12 @@ def test_getcl_lmax_0(
     )
     expected = xpb.asarray([max(random_i, random_j)], dtype=xpb.float64)
     assert result.shape[0] == 1
-    compare.assert_allclose(result, expected)
+    xpx.testing.assert_close(result, expected)
 
 
 @pytest.mark.unstable
 def test_getcl_lmax_larger_than_cls(
     benchmark: BenchmarkFixture,
-    compare: type[Compare],
     xpb: ModuleType,
 ) -> None:
     """Benchmarks for glass.getcl with lmax larger than the length of cl."""
@@ -277,4 +274,4 @@ def test_getcl_lmax_larger_than_cls(
     )
     expected = xpb.zeros((lmax - 1,), dtype=xpb.float64)
     assert result.shape[0] == lmax + 1
-    compare.assert_allclose(result[2:], expected)
+    xpx.testing.assert_close(result[2:], expected)

tests/benchmarks/test_harmonics.py

@@ -4,6 +4,8 @@
 
 import pytest
 
+import array_api_extra as xpx
+
 glass_harmonics = pytest.importorskip(
     "glass.harmonics",
     reason="tests require glass.harmonics",
@@ -15,13 +17,10 @@
 
     from pytest_benchmark.fixture import BenchmarkFixture
 
-    from tests.fixtures.helper_classes import Compare
-
 
 @pytest.mark.unstable
 def test_multalm(
     benchmark: BenchmarkFixture,
-    compare: type[Compare],
     xpb: ModuleType,
 ) -> None:
     """Benchmarks for glass.harmonics.multalm."""
@@ -32,7 +31,7 @@ def test_multalm(
 
     result = benchmark(glass_harmonics.multalm, alm, bl)
 
-    compare.assert_allclose(
+    xpx.testing.assert_close(
         result[:5],
         xpb.asarray([scale_factor * x for x in [0.0, 3.0, 6.0, 9.0, 12.0]]),
     )

tests/benchmarks/test_lensing.py

@@ -4,6 +4,8 @@
 
 import pytest
 
+import array_api_extra as xpx
+
 import glass
 
 if TYPE_CHECKING:
@@ -14,13 +16,11 @@
 
     from glass._types import FloatArray, UnifiedGenerator
     from glass.cosmology import Cosmology
-    from tests.fixtures.helper_classes import Compare
 
 
 @pytest.mark.stable
 def test_multi_plane_matrix(
     benchmark: BenchmarkFixture,
-    compare: type[Compare],
     cosmo: Cosmology,
     urngb: UnifiedGenerator,
     xpb: ModuleType,
@@ -38,8 +38,8 @@ def test_multi_plane_matrix(
     mat = glass.multi_plane_matrix(shells, cosmo)
     deltas = urngb.random((len(shells), 10))
 
-    compare.assert_array_equal(mat, xpb.tril(mat))
-    compare.assert_array_equal(xpb.triu(mat, k=1), 0)
+    xpx.testing.assert_equal(mat, xpb.tril(mat))
+    xpx.testing.assert_equal(xpb.triu(mat, k=1), xpb.asarray(0.0), check_shape=False)
 
     def setup_shells_and_deltas() -> tuple[
         tuple[
@@ -79,7 +79,6 @@ def multi_plane_matrix_add_window(
 @pytest.mark.stable
 def test_multi_plane_weights(
     benchmark: BenchmarkFixture,
-    compare: type[Compare],
     cosmo: Cosmology,
     urngb: UnifiedGenerator,
     xpb: ModuleType,
@@ -100,8 +99,8 @@ def test_multi_plane_weights(
 
     w_out = glass.multi_plane_weights(w_in, shells, cosmo)
 
-    compare.assert_array_equal(w_out, xpb.triu(w_out, 1))
-    compare.assert_array_equal(xpb.tril(w_out), 0)
+    xpx.testing.assert_equal(w_out, xpb.triu(w_out, 1))
+    xpx.testing.assert_equal(xpb.tril(w_out), xpb.asarray(0.0), check_shape=False)
 
     def setup_shells_deltas_and_weights() -> tuple[
         tuple[

tests/benchmarks/test_shells.py

@@ -4,15 +4,15 @@
 
 import pytest
 
+import array_api_extra as xpx
+
 import glass
 
 if TYPE_CHECKING:
     from types import ModuleType
 
     from pytest_benchmark.fixture import BenchmarkFixture
 
-    from tests.fixtures.helper_classes import Compare
-
 
 @pytest.mark.unstable
 def test_radialwindow(
@@ -23,11 +23,11 @@ def test_radialwindow(
     """Benchmark for shells.RadialWindow."""
     # check zeff is computed when not provided
     arr_length = 100_000
-    expected_zeff = 66_666.0
+    expected_zeff = xpb.asarray(66_666.0)
 
     wa = xpb.arange(arr_length)
     za = xpb.arange(arr_length)
 
     w = benchmark(glass.RadialWindow, za, wa)
 
-    compare.assert_allclose(w.zeff, expected_zeff)
+    xpx.testing.assert_close(w.zeff, expected_zeff)

tests/core/grf/test_solver.py

@@ -5,11 +5,12 @@
 import numpy as np
 import pytest
 
+import array_api_extra as xpx
+
 import glass.grf
 
 if TYPE_CHECKING:
     from glass._types import FloatArray
-    from tests.fixtures.helper_classes import Compare
 
 
 @pytest.fixture(scope="session")
@@ -21,7 +22,6 @@ def cl() -> FloatArray:
 
 def test_one_transformation(
     cl: FloatArray,
-    compare: type[Compare],
     rng: np.random.Generator,
 ) -> None:
     lam = rng.random()
@@ -30,7 +30,7 @@ def test_one_transformation(
     gl1, _, _ = glass.grf.solve(cl, t)
     gl2, _, _ = glass.grf.solve(cl, t, t)
 
-    compare.assert_array_equal(gl1, gl2)
+    xpx.testing.assert_equal(gl1, gl2)
 
 
 def test_pad(
@@ -51,7 +51,6 @@ def test_pad(
 
 def test_initial(
     cl: FloatArray,
-    compare: type[Compare],
     rng: np.random.Generator,
 ) -> None:
     lam = rng.random()
@@ -62,24 +61,20 @@ def test_initial(
     gl1, _, _ = glass.grf.solve(cl, t)
     gl2, _, _ = glass.grf.solve(cl, t, initial=gl)
 
-    compare.assert_array_equal(gl1, gl2)
+    xpx.testing.assert_equal(gl1, gl2)
 
 
-def test_no_iterations(
-    cl: FloatArray,
-    compare: type[Compare],
-) -> None:
+def test_no_iterations(cl: FloatArray) -> None:
     t = glass.grf.Lognormal()
 
     gl1 = glass.grf.compute(cl, t)
     gl2, _, _ = glass.grf.solve(cl, t, maxiter=0)
 
-    compare.assert_array_equal(gl1, gl2)
+    xpx.testing.assert_equal(gl1, gl2)
 
 
 def test_lognormal(
     cl: FloatArray,
-    compare: type[Compare],
     rng: np.random.Generator,
 ) -> None:
     t1 = glass.grf.Lognormal()
@@ -93,17 +88,16 @@ def test_lognormal(
 
     assert info > 0
 
-    compare.assert_allclose(cl_[1 : cl.shape[0]], cl[1:], atol=0.0, rtol=cltol)
+    xpx.testing.assert_close(cl_[1 : cl.shape[0]], cl[1:], atol=0.0, rtol=cltol)
 
     gl_ = glass.grf.compute(cl_, t1, t2)
 
     assert gl[0] == gl0
-    compare.assert_allclose(gl_[1 : gl.shape[0]], gl[1:])
+    xpx.testing.assert_close(gl_[1 : gl.shape[0]], gl[1:])
 
 
 def test_monopole(
     cl: FloatArray,
-    compare: type[Compare],
     rng: np.random.Generator,
 ) -> None:
     t = glass.grf.Lognormal()
@@ -114,10 +108,10 @@ def test_monopole(
     gl, cl_out, _ = glass.grf.solve(cl, t, monopole=None, gltol=1e-8)
 
     assert gl[0] != 0.0
-    compare.assert_allclose(cl_out[0], cl[0])
+    xpx.testing.assert_close(cl_out[0], cl[0])
 
     gl, cl_out, _ = glass.grf.solve(cl, t, monopole=gl0, gltol=1e-8)
 
     assert gl[0] == gl0
     with pytest.raises(AssertionError, match="Not equal to tolerance"):
-        compare.assert_allclose(cl_out[0], cl[0])
+        xpx.testing.assert_close(cl_out[0], cl[0])