From 95ef6a4ab2aa2eb508d8c17712cd873f0ce36adf Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 29 May 2026 08:59:43 +0100 Subject: [PATCH] Honour PYAUTO_TEST_MODE in LOSSampler to fix los_halos simulator timeouts LOSSampler.galaxies_from now caps the line-of-sight halo population to a few per plane and loosens the negative-kappa quad integration when PYAUTO_TEST_MODE is active. This collapses the los_halos simulator runtime under the workspace integration runner (simulator.py 320s timeout -> 28s; simulator_jax.py timeout -> 45s) while keeping the full sampling and ray-tracing code paths exercised. Non-test-mode behaviour is bit-identical (still samples ~1100 halos with full-accuracy kappa). Co-Authored-By: Claude Opus 4.8 (1M context) --- autolens/lens/los.py | 87 +++++++++++++++++++++----- test_autolens/lens/test_los.py | 110 +++++++++++++++++++++++++++++++++ 2 files changed, 180 insertions(+), 17 deletions(-) create mode 100644 test_autolens/lens/test_los.py diff --git a/autolens/lens/los.py b/autolens/lens/los.py index 162045066..70d71522e 100644 --- a/autolens/lens/los.py +++ b/autolens/lens/los.py @@ -12,14 +12,24 @@ multi-plane :class:`autolens.Tracer`. """ +import warnings + import numpy as np -from scipy.integrate import quad +from scipy.integrate import quad, IntegrationWarning from scipy.interpolate import interp1d from typing import List, Optional, Tuple import autogalaxy as ag from autogalaxy.cosmology import Planck15 +from autoconf.test_mode import is_test_mode + +# Number of LOS halos retained per plane when ``PYAUTO_TEST_MODE`` is active. +# Capping the population keeps the multi-plane ray-tracing and per-galaxy +# plotting paths exercised (so regressions still surface) while collapsing the +# downstream cost from ~1100 halos to a few dozen. See ``LOSSampler.galaxies_from``. +_TEST_MODE_MAX_HALOS_PER_PLANE = 3 + def comoving_distance_mpc_from(z, cosmology): """ @@ -241,6 +251,8 @@ def negative_kappa_from( truncation_factor, c_scatter, cosmology, + quad_limit=50, + quad_epsrel=1.49e-8, ): """ Compute the negative convergence sheet for a single LOS plane. @@ -278,6 +290,16 @@ def negative_kappa_from( Log-normal scatter in concentration (sigma in dex, e.g. 0.15). cosmology A ``LensingCosmology`` instance. + quad_limit + Maximum number of adaptive subintervals for both the inner + (concentration) and outer (mass) ``scipy.integrate.quad`` calls. + Defaults to scipy's own default of ``50``. ``LOSSampler.galaxies_from`` + lowers this under ``PYAUTO_TEST_MODE`` to make the double integral cheap + while still exercising the full integrand (the inner ``fsolve`` is the + dominant cost, so fewer subintervals is a ~50x speed-up). + quad_epsrel + Relative error tolerance passed to both ``quad`` calls. Defaults to + scipy's own default of ``1.49e-8``; loosened under test mode. Returns ------- @@ -316,13 +338,20 @@ def _integrand_mass(m): lgc_hi = lgc_centre + 4.0 * c_scatter c_integral = quad( - _integrand_concentration, lgc_lo, lgc_hi, args=(m, lgc_centre) + _integrand_concentration, + lgc_lo, + lgc_hi, + args=(m, lgc_centre), + limit=quad_limit, + epsrel=quad_epsrel, )[0] dndm = 10 ** B_mf * m ** A_mf return dndm * m * c_integral - mass_integral = quad(_integrand_mass, m_min, m_max)[0] + mass_integral = quad( + _integrand_mass, m_min, m_max, limit=quad_limit, epsrel=quad_epsrel + )[0] kappa = mass_integral * comoving_volume_per_arcsec2 / sigma_cr_mpc2 @@ -601,6 +630,17 @@ def galaxies_from(self) -> List[ag.Galaxy]: cosmology = self.cosmology rng = np.random.RandomState(self.seed) + # ``PYAUTO_TEST_MODE`` (integration tests / workspace smoke runs) makes + # the full LOS population prohibitively slow: a science run samples + # ~1100 halos (driving multi-plane ray tracing to ~90s) and the + # per-plane negative-kappa double integral costs ~3.8s/plane. Under test + # mode we cap the halos per plane and loosen the kappa integral, which + # keeps both code paths exercised while collapsing the runtime so the + # los_halos simulators finish well under the per-script timeout cap. + test_mode = is_test_mode() + quad_limit = 1 if test_mode else 50 + quad_epsrel = 0.1 if test_mode else 1.49e-8 + boundaries, centres = los_planes_from( z_lens=self.z_lens, z_source=self.z_source, @@ -682,6 +722,9 @@ def galaxies_from(self) -> List[ag.Galaxy]: ) n_halos = rng.poisson(n_bar) + if test_mode: + n_halos = min(n_halos, _TEST_MODE_MAX_HALOS_PER_PLANE) + if n_halos > 0: masses = sample_halo_masses( n=n_halos, @@ -718,20 +761,30 @@ def galaxies_from(self) -> List[ag.Galaxy]: ag.Galaxy(redshift=z_cen, mass=halo) ) - kappa_neg = negative_kappa_from( - z_centre=z_cen, - comoving_volume_per_arcsec2=vol_depth, - A_mf=mf_coeffs[i, 0], - B_mf=mf_coeffs[i, 1], - A_mc=mc_coeffs[i, 0], - B_mc=mc_coeffs[i, 1], - m_min=self.m_min, - m_max=self.m_max, - z_source=self.z_source, - truncation_factor=self.truncation_factor, - c_scatter=self.c_scatter, - cosmology=cosmology, - ) + with warnings.catch_warnings(): + # Under test mode the deliberately low ``quad_limit`` makes + # scipy emit a (harmless, expected) max-subdivisions warning per + # integral; silence it so smoke-run output stays clean. Full + # accuracy runs (quad_limit=50) never trip it. + if test_mode: + warnings.simplefilter("ignore", IntegrationWarning) + + kappa_neg = negative_kappa_from( + z_centre=z_cen, + comoving_volume_per_arcsec2=vol_depth, + A_mf=mf_coeffs[i, 0], + B_mf=mf_coeffs[i, 1], + A_mc=mc_coeffs[i, 0], + B_mc=mc_coeffs[i, 1], + m_min=self.m_min, + m_max=self.m_max, + z_source=self.z_source, + truncation_factor=self.truncation_factor, + c_scatter=self.c_scatter, + cosmology=cosmology, + quad_limit=quad_limit, + quad_epsrel=quad_epsrel, + ) galaxies.append( ag.Galaxy( redshift=z_cen, diff --git a/test_autolens/lens/test_los.py b/test_autolens/lens/test_los.py new file mode 100644 index 000000000..85176e59c --- /dev/null +++ b/test_autolens/lens/test_los.py @@ -0,0 +1,110 @@ +import numpy as np +import pytest + +import autolens as al +from autogalaxy.cosmology import Planck15 +from autolens.lens import los + + +def _approx_coefficients(n_planes): + """ + Approximate per-plane mass-function and mass-concentration coefficients, + matching the pre-computed fallback used by the los_halos workspace + simulators (avoids the optional ``hmf`` / ``colossus`` dependencies). + """ + mf = np.tile([-1.9, 8.0], (n_planes, 1)) + mc = np.tile([-3.0, 40.0], (n_planes, 1)) + return mf, mc + + +def test__negative_kappa_from__loose_quad_matches_reference_and_is_negative(): + """ + The ``quad_limit`` / ``quad_epsrel`` knobs that test mode lowers must thread + through to both ``quad`` calls and still produce a finite, negative kappa. + A coarse integration (limit=1) should agree with a finer one to a few + percent — the value is unused in test mode, so loose accuracy is fine. + """ + cosmology = Planck15() + _, centres = los.los_planes_from( + z_lens=0.5, z_source=1.0, planes_before_lens=4, planes_after_lens=4 + ) + mf, mc = _approx_coefficients(len(centres)) + + kwargs = dict( + z_centre=centres[0], + comoving_volume_per_arcsec2=1.0, + A_mf=mf[0, 0], + B_mf=mf[0, 1], + A_mc=mc[0, 0], + B_mc=mc[0, 1], + m_min=1e7, + m_max=1e10, + z_source=1.0, + truncation_factor=100.0, + c_scatter=0.15, + cosmology=cosmology, + ) + + reference = los.negative_kappa_from(quad_limit=10, quad_epsrel=1e-3, **kwargs) + coarse = los.negative_kappa_from(quad_limit=1, quad_epsrel=1e-1, **kwargs) + + assert reference < 0.0 + assert coarse < 0.0 + assert coarse == pytest.approx(reference, rel=0.05) + + +def test__galaxies_from__test_mode_caps_halos_per_plane(monkeypatch): + """ + Under ``PYAUTO_TEST_MODE`` ``galaxies_from`` must cap the halo population to + a handful per plane (so the downstream multi-plane ray tracing stays cheap) + while still emitting one negative-kappa ``MassSheet`` galaxy per plane. + """ + monkeypatch.setenv("PYAUTO_TEST_MODE", "2") + + cosmology = Planck15() + _, centres = los.los_planes_from( + z_lens=0.5, z_source=1.0, planes_before_lens=4, planes_after_lens=4 + ) + n_planes = len(centres) + mf, mc = _approx_coefficients(n_planes) + + sampler = los.LOSSampler( + z_lens=0.5, + z_source=1.0, + planes_before_lens=4, + planes_after_lens=4, + m_min=1e7, + m_max=1e10, + cone_radius_arcsec=5.0, + c_scatter=0.15, + truncation_factor=100.0, + cosmology=cosmology, + mass_function_coefficients=mf, + mass_concentration_coefficients=mc, + seed=42, + ) + + galaxies = sampler.galaxies_from() + + halos = [ + g + for g in galaxies + if hasattr(g, "mass") and isinstance(g.mass, al.mp.NFWTruncatedSph) + ] + sheets = [ + g + for g in galaxies + if hasattr(g, "mass_sheet") and isinstance(g.mass_sheet, al.mp.MassSheet) + ] + + # One negative-kappa sheet per plane, all with negative convergence. + assert len(sheets) == n_planes + assert all(g.mass_sheet.kappa < 0.0 for g in sheets) + + # Halos are capped to at most three per plane (grouped by plane redshift). + counts = {} + for g in halos: + counts[g.redshift] = counts.get(g.redshift, 0) + 1 + + assert len(counts) > 0 + assert all(count <= 3 for count in counts.values())