update benchmark suite

tests/benchmarks/test_actionangle_benchmark.py

@@ -1,3 +1,5 @@
+import os
+
 import astropy.units as u
 import numpy as np
 import pytest
@@ -35,7 +37,14 @@ def test_staeckel_fudge(potential, benchmark, n_points):
 
 @pytest.mark.parametrize("n_steps", [1_000, 10_000])
 @pytest.mark.parametrize("n_orbits", [1, 10])
-def test_o2gf(potential, benchmark, n_orbits, n_steps):
+def test_o2gf(request, potential, benchmark, n_orbits, n_steps):
+    # Codspeed takes an hour for 10 orbits; stick to 1.
+    codspeed = (
+        request.config.getoption("--codspeed", False) or "CODSPEED_ENV" in os.environ
+    )
+    if codspeed and n_orbits > 1:
+        pytest.skip("multiple orbits are only timed for the scaling benchmarks")
+
     # integrate outside the timed region: integration has its own benchmarks
     orbits = potential.integrate_orbit(
         make_w0(potential, n_orbits), dt=1 * u.Myr, n_steps=n_steps

tests/benchmarks/test_integrate_benchmark.py

@@ -1,4 +1,5 @@
 import astropy.units as u
+import numpy as np
 import pytest
 
 import gala.dynamics as gd
@@ -10,21 +11,45 @@ def potential():
     return gp.HernquistPotential(m=1e11, c=5, units="galactic")
 
 
+def make_w0(v_y, n_orbits, n_steps):
+    if n_orbits * n_steps > 1_000_000:
+        pytest.skip("too slow under CodSpeed's simulation")
+    xyz = np.zeros((3, n_orbits))
+    xyz[0] = 10.0
+    vxyz = np.zeros((3, n_orbits))
+    vxyz[1] = v_y
+    return gd.PhaseSpacePosition(xyz * u.kpc, vxyz * u.km / u.s)
+
+
+@pytest.fixture
+def circular_w0(n_orbits, n_steps):
+    return make_w0(138.25768647, n_orbits, n_steps)
+
+
+@pytest.fixture
+def eccentric_w0(n_orbits, n_steps):
+    return make_w0(0.25 * 138.25768647, n_orbits, n_steps)
+
+
 @pytest.mark.benchmark
 @pytest.mark.parametrize("Integrator", ["dop853", "leapfrog", "ruth4"])
 @pytest.mark.parametrize("n_steps", [1, 10, 1_000, 100_000])
 @pytest.mark.parametrize("n_orbits", [1, 10, 1_000, 100_000])
-def test_integrate_circular_orbit(potential, Integrator, n_steps, n_orbits):
-    w0 = gd.PhaseSpacePosition([10.0, 0, 0] * u.kpc, [0, 138.25768647, 0] * u.km / u.s)
-    potential.integrate_orbit(w0, dt=1 * u.Myr, n_steps=n_steps, Integrator=Integrator)
+def test_integrate_circular_orbit(
+    potential, circular_w0, Integrator, n_steps, n_orbits
+):
+    potential.integrate_orbit(
+        circular_w0, dt=1 * u.Myr, n_steps=n_steps, Integrator=Integrator
+    )
 
 
 @pytest.mark.benchmark
 @pytest.mark.parametrize("Integrator", ["dop853", "leapfrog", "ruth4"])
 @pytest.mark.parametrize("n_steps", [1, 10, 1_000, 100_000])
 @pytest.mark.parametrize("n_orbits", [1, 10, 1_000, 100_000])
-def test_integrate_eccentric_orbit(potential, Integrator, n_steps, n_orbits):
-    w0 = gd.PhaseSpacePosition(
-        [10.0, 0, 0] * u.kpc, [0, 0.25 * 138.25768647, 0] * u.km / u.s
+def test_integrate_eccentric_orbit(
+    potential, eccentric_w0, Integrator, n_steps, n_orbits
+):
+    potential.integrate_orbit(
+        eccentric_w0, dt=1 * u.Myr, n_steps=n_steps, Integrator=Integrator
     )
-    potential.integrate_orbit(w0, dt=1 * u.Myr, n_steps=n_steps, Integrator=Integrator)

tests/benchmarks/test_potentials_benchmark.py

@@ -6,6 +6,7 @@
 import pytest
 from gala._cconfig import EXP_ENABLED, GSL_ENABLED
 
+import gala.dynamics as gd
 import gala.potential as gp
 from gala.units import SimulationUnitSystem, galactic
 
@@ -28,27 +29,28 @@ class BenchmarkPotentialBase:
     tested (e.g., for expensive potentials that are too slow at large n_points).
     """
 
-    @pytest.fixture(params=[1, 10, 1_000])
+    @pytest.fixture(params=[1, 1_000, 10_000])
     def n_points(self, request):
         return request.param
 
-    @pytest.fixture(scope="class")
-    def rng(self):
-        return np.random.default_rng(42)
+    @pytest.fixture
+    def q(self, potential, n_points):
+        # A fixture, so that generating the positions is not timed
+        return np.random.default_rng(42).normal(0, 10, size=(potential.ndim, n_points))
 
     @pytest.mark.benchmark
-    def test_energy(self, potential, n_points, rng):
-        potential.energy(rng.normal(0, 10, size=(potential.ndim, n_points)))
+    def test_energy(self, potential, q):
+        potential.energy(q)
 
     @pytest.mark.benchmark
-    def test_gradient(self, potential, n_points, rng):
-        potential.gradient(rng.normal(0, 10, size=(potential.ndim, n_points)))
+    def test_gradient(self, potential, q):
+        potential.gradient(q)
 
     @pytest.mark.benchmark
-    def test_density(self, potential, n_points, rng):
+    def test_density(self, potential, q):
         if type(potential).__name__ in NO_DENSITY:
             pytest.skip("density not implemented for this potential")  # type: ignore[call-arg]
-        potential.density(rng.normal(0, 10, size=(potential.ndim, n_points)))
+        potential.density(q)
 
 
 # ============================================================================
@@ -68,6 +70,18 @@ class TestCanonicalPotentialsBenchmark(BenchmarkPotentialBase):
     pass
 
 
+@pytest.fixture
+def positions(n_points):
+    return np.random.default_rng(42).normal(0, 10, size=(3, n_points))
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize("n_points", [1, 1_000, 10_000])
+def test_hessian(positions):
+    # Just one representative potential
+    CANONICAL["HernquistPotential"].hessian(positions)
+
+
 # ============================================================================
 # Special (file-based potentials with expensive setup)
 
@@ -139,16 +153,16 @@ def potential(self):
         )
 
     @pytest.mark.benchmark
-    def test_energy(self, potential, n_points, rng):
-        potential.energy(rng.normal(0, 10, size=(potential.ndim, n_points)), t=2.5)
+    def test_energy(self, potential, q):
+        potential.energy(q, t=2.5)
 
     @pytest.mark.benchmark
-    def test_gradient(self, potential, n_points, rng):
-        potential.gradient(rng.normal(0, 10, size=(potential.ndim, n_points)), t=2.5)
+    def test_gradient(self, potential, q):
+        potential.gradient(q, t=2.5)
 
     @pytest.mark.benchmark
-    def test_density(self, potential, n_points, rng):
-        potential.density(rng.normal(0, 10, size=(potential.ndim, n_points)), t=2.5)
+    def test_density(self, potential, q):
+        potential.density(q, t=2.5)
 
 
 @pytest.mark.skipif(not EXP_ENABLED, reason="requires EXP")
@@ -187,3 +201,40 @@ def potential(self):
             coef_file=EXP_SINGLE_COEF_FILE,
             units=exp_units,
         )
+
+
+# ============================================================================
+# Input types other than plain arrays, which are converted before evaluation
+
+
+@pytest.fixture
+def positions_kpc(positions):
+    return positions * u.kpc
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize("n_points", [1, 10_000])
+def test_energy_quantity(positions_kpc):
+    CANONICAL["HernquistPotential"].energy(positions_kpc)
+
+
+@pytest.fixture(scope="module")
+def orbit():
+    # 1,000 orbits x 1,000 times: large enough that copies of the orbit cost more
+    # than the fixed overhead of a call
+    rng = np.random.default_rng(42)
+    w0 = gd.PhaseSpacePosition(
+        pos=rng.normal(0, 10, size=(3, 1000)) * u.kpc,
+        vel=rng.normal(0, 100, size=(3, 1000)) * u.km / u.s,
+    )
+    return CANONICAL["HernquistPotential"].integrate_orbit(w0, dt=1.0, n_steps=999)
+
+
+@pytest.mark.benchmark
+def test_energy_orbit(orbit):
+    CANONICAL["HernquistPotential"].energy(orbit)
+
+
+@pytest.mark.benchmark
+def test_orbit_energy_method(orbit):
+    orbit.energy()