diff --git a/tests/benchmarks/test_actionangle_benchmark.py b/tests/benchmarks/test_actionangle_benchmark.py index b7e6b8eab..342b1c4f3 100644 --- a/tests/benchmarks/test_actionangle_benchmark.py +++ b/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 diff --git a/tests/benchmarks/test_integrate_benchmark.py b/tests/benchmarks/test_integrate_benchmark.py index a8aaed410..d0ff18d8f 100644 --- a/tests/benchmarks/test_integrate_benchmark.py +++ b/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) diff --git a/tests/benchmarks/test_potentials_benchmark.py b/tests/benchmarks/test_potentials_benchmark.py index c3181a156..baa0635e0 100644 --- a/tests/benchmarks/test_potentials_benchmark.py +++ b/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()