test: redo benchmarks to just cover jax and add timing comps

Author: beckermr

jax_galsim/spergel.py

@@ -267,34 +267,34 @@ def scale_radius(self):
     def _r0(self):
         return self.scale_radius
 
-    @property
+    @lazy_property
     def _inv_r0(self):
         return 1.0 / self._r0
 
-    @property
+    @lazy_property
     def _r0_sq(self):
         return self._r0 * self._r0
 
-    @property
+    @lazy_property
     def _inv_r0_sq(self):
         return self._inv_r0 * self._inv_r0
 
-    @property
+    @lazy_property
     @implements(_galsim.spergel.Spergel.half_light_radius)
     def half_light_radius(self):
         return self._r0 * calculateFluxRadius(0.5, self.nu)
 
-    @property
+    @lazy_property
     def _shootxnorm(self):
         """Normalization for photon shooting"""
         return 1.0 / (2.0 * jnp.pi * jnp.power(2.0, self.nu) * _gammap1(self.nu))
 
-    @property
+    @lazy_property
     def _xnorm(self):
         """Normalization of xValue"""
         return self._shootxnorm * self.flux * self._inv_r0_sq
 
-    @property
+    @lazy_property
     def _xnorm0(self):
         """return z^nu K_nu(z) for z=0"""
         return jax.lax.select(
@@ -338,21 +338,21 @@ def __str__(self):
         s += ")"
         return s
 
-    @property
+    @lazy_property
     def _maxk(self):
         """(1+ (k r0)^2)^(-1-nu) = maxk_threshold"""
         res = jnp.power(self.gsparams.maxk_threshold, -1.0 / (1.0 + self.nu)) - 1.0
         return jnp.sqrt(res) / self._r0
 
-    @property
+    @lazy_property
     def _stepk(self):
         R = calculateFluxRadius(1.0 - self.gsparams.folding_threshold, self.nu)
         R *= self._r0
         # Go to at least 5*hlr
         R = jnp.maximum(R, self.gsparams.stepk_minimum_hlr * self.half_light_radius)
         return jnp.pi / R
 
-    @property
+    @lazy_property
     def _max_sb(self):
         # from SBSpergelImpl.h
         return jnp.abs(self._xnorm) * self._xnorm0

tests/jax/test_benchmarks.py

@@ -179,18 +179,12 @@ def test_benchmark_spergel_conv(benchmark, kind):
     dt = _run_benchmarks(
         benchmark, kind, lambda: _run_spergel_bench_conv_jit().block_until_ready()
     )
-    print(f"jax-galsim time: {dt:0.4g} ms", end=" ")
-
-
-@pytest.mark.parametrize("kind", ["compile", "run"])
-def test_benchmark_spergel_conv_galsim(benchmark, kind):
-    dt = _run_benchmarks(benchmark, kind, lambda: _run_spergel_bench_conv(_galsim))
-    print(f"galsim time: {dt:0.4g} ms", end=" ")
+    print(f"time: {dt:0.4g} ms", end=" ")
 
 
 def _run_spergel_bench_xvalue(gsmod):
     obj = gsmod.Spergel(nu=-0.6, scale_radius=5)
-    return obj.drawImage(nx=50, ny=50, scale=0.2, method="no_pixel").array
+    return obj.drawImage(nx=1024, ny=1204, scale=0.05, method="no_pixel").array
 
 
 _run_spergel_bench_xvalue_jit = jax.jit(partial(_run_spergel_bench_xvalue, jgs))
@@ -201,18 +195,12 @@ def test_benchmark_spergel_xvalue(benchmark, kind):
     dt = _run_benchmarks(
         benchmark, kind, lambda: _run_spergel_bench_xvalue_jit().block_until_ready()
     )
-    print(f"jax-galsim time: {dt:0.4g} ms", end=" ")
-
-
-@pytest.mark.parametrize("kind", ["compile", "run"])
-def test_benchmark_spergel_xvalue_galsim(benchmark, kind):
-    dt = _run_benchmarks(benchmark, kind, lambda: _run_spergel_bench_xvalue(_galsim))
-    print(f"galsim time: {dt:0.4g} ms", end=" ")
+    print(f"time: {dt:0.4g} ms", end=" ")
 
 
 def _run_spergel_bench_kvalue(gsmod):
     obj = gsmod.Spergel(nu=-0.6, scale_radius=5)
-    return obj.drawKImage(nx=50, ny=50, scale=0.2).array
+    return obj.drawKImage(nx=1024, ny=1204, scale=0.05).array
 
 
 _run_spergel_bench_kvalue_jit = jax.jit(partial(_run_spergel_bench_kvalue, jgs))
@@ -223,10 +211,4 @@ def test_benchmark_spergel_kvalue(benchmark, kind):
     dt = _run_benchmarks(
         benchmark, kind, lambda: _run_spergel_bench_kvalue_jit().block_until_ready()
     )
-    print(f"jax-galsim time: {dt:0.4g} ms", end=" ")
-
-
-@pytest.mark.parametrize("kind", ["compile", "run"])
-def test_benchmark_spergel_kvalue_galsim(benchmark, kind):
-    dt = _run_benchmarks(benchmark, kind, lambda: _run_spergel_bench_kvalue(_galsim))
-    print(f"galsim time: {dt:0.4g} ms", end=" ")
+    print(f"time: {dt:0.4g} ms", end=" ")

tests/jax/test_spergel_comp_galsim.py

@@ -1,8 +1,19 @@
+import galsim as _galsim
 import galsim as gs
+import jax
 import numpy as np
 import pytest
+from test_benchmarks import (
+    _run_spergel_bench_conv,
+    _run_spergel_bench_conv_jit,
+    _run_spergel_bench_kvalue,
+    _run_spergel_bench_kvalue_jit,
+    _run_spergel_bench_xvalue,
+    _run_spergel_bench_xvalue_jit,
+)
 
 import jax_galsim as jgs
+from jax_galsim.core.testing import time_code_block
 
 
 @pytest.mark.parametrize(
@@ -78,3 +89,69 @@ def test_spergel_comp_galsim_xvalue(nu, scale_radius, x, y):
     s_gs = gs.Spergel(nu=nu, scale_radius=scale_radius)
 
     np.testing.assert_allclose(s_jgs.xValue(x, y), s_gs.xValue(x, y), rtol=1e-5)
+
+
+def _run_time_test(kind, func):
+    if kind == "compile":
+
+        def _run():
+            jax.clear_caches()
+            func()
+
+    elif kind == "run":
+        # run once to compile
+        func()
+
+        def _run():
+            func()
+
+    else:
+        raise ValueError(f"kind={kind} not recognized")
+
+    tot_time = 0
+    for _ in range(3):
+        with time_code_block(quiet=True) as tr:
+            _run()
+        tot_time += tr.dt
+
+    return tot_time / 3
+
+
+@pytest.mark.parametrize("kind", ["compile", "run"])
+def test_spergel_comp_galsim_perf_conv(benchmark, kind):
+    dt = _run_time_test(kind, lambda: _run_spergel_bench_conv_jit().block_until_ready())
+    print(f"\njax-galsim time: {dt:0.4g} ms")
+
+    dt = _run_time_test(
+        kind,
+        lambda: _run_spergel_bench_conv(_galsim),
+    )
+    print(f"    galsim time: {dt:0.4g} ms")
+
+
+@pytest.mark.parametrize("kind", ["compile", "run"])
+def test_spergel_comp_galsim_perf_kvalue(benchmark, kind):
+    dt = _run_time_test(
+        kind, lambda: _run_spergel_bench_kvalue_jit().block_until_ready()
+    )
+    print(f"\njax-galsim time: {dt:0.4g} ms")
+
+    dt = _run_time_test(
+        kind,
+        lambda: _run_spergel_bench_kvalue(_galsim),
+    )
+    print(f"    galsim time: {dt:0.4g} ms")
+
+
+@pytest.mark.parametrize("kind", ["compile", "run"])
+def test_spergel_comp_galsim_perf_xvalue(benchmark, kind):
+    dt = _run_time_test(
+        kind, lambda: _run_spergel_bench_xvalue_jit().block_until_ready()
+    )
+    print(f"\njax-galsim time: {dt:0.4g} ms")
+
+    dt = _run_time_test(
+        kind,
+        lambda: _run_spergel_bench_xvalue(_galsim),
+    )
+    print(f"    galsim time: {dt:0.4g} ms")