From 413c6eb5369b7a75f5388fbb2a05b0bddb2898e9 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sat, 11 Jul 2026 11:45:25 +0100 Subject: [PATCH] test: skip jax-backed tests when jax is absent (Python matrix) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit These tests exercise jax-only code paths (vmapped profiles / blackjax NUTS / nufft sparse operator / use_jax=True), so they need jax installed to run — it ships via the [optional] extras and is present in the per-repo CI (3.12/3.13). The PyAutoBuild Python Version Matrix installs the NumPy-only stack (and jax can't install on 3.9 anyway), so guard the 5 affected test(s) with a skipif on jax availability rather than letting them ModuleNotFoundError. Numpy-only tests in the same files are unaffected. Co-Authored-By: Claude Opus 4.8 --- .../model/test_analysis_interferometer.py | 365 +++++++++--------- test_autolens/lens/test_substructure_util.py | 12 + 2 files changed, 201 insertions(+), 176 deletions(-) diff --git a/test_autolens/interferometer/model/test_analysis_interferometer.py b/test_autolens/interferometer/model/test_analysis_interferometer.py index 5b796e869..c1bd2f4f4 100644 --- a/test_autolens/interferometer/model/test_analysis_interferometer.py +++ b/test_autolens/interferometer/model/test_analysis_interferometer.py @@ -1,176 +1,189 @@ -from pathlib import Path -import pytest - -import autofit as af -import autoarray as aa -import autolens as al -from autolens import exc - -from autolens.interferometer.model.result import ResultInterferometer - -directory = Path(__file__).resolve().parent - - -def test__make_result__result_interferometer_is_returned(interferometer_7): - model = af.Collection(galaxies=af.Collection(galaxy_0=al.Galaxy(redshift=0.5))) - - analysis = al.AnalysisInterferometer(dataset=interferometer_7, use_jax=False) - - search = al.m.MockSearch(name="test_search") - - result = search.fit(model=model, analysis=analysis) - - assert isinstance(result, ResultInterferometer) - - -def test__figure_of_merit__matches_correct_fit_given_galaxy_profiles(interferometer_7): - lens_galaxy = al.Galaxy(redshift=0.5, light=al.lp.Sersic(intensity=0.1)) - - model = af.Collection(galaxies=af.Collection(lens=lens_galaxy)) - - analysis = al.AnalysisInterferometer(dataset=interferometer_7, use_jax=False) - - instance = model.instance_from_unit_vector([]) - analysis_log_likelihood = analysis.log_likelihood_function(instance=instance) - - tracer = analysis.tracer_via_instance_from(instance=instance) - - fit = al.FitInterferometer(dataset=interferometer_7, tracer=tracer) - - assert fit.log_likelihood == analysis_log_likelihood - - -def test__positions__likelihood_overwrite__changes_likelihood( - interferometer_7, mask_2d_7x7 -): - lens = al.Galaxy(redshift=0.5, mass=al.mp.IsothermalSph(centre=(0.05, 0.05))) - source = al.Galaxy(redshift=1.0, light=al.lp.SersicSph(centre=(0.05, 0.05))) - - model = af.Collection(galaxies=af.Collection(lens=lens, source=source)) - - analysis = al.AnalysisInterferometer(dataset=interferometer_7, use_jax=False) - - instance = model.instance_from_unit_vector([]) - analysis_log_likelihood = analysis.log_likelihood_function(instance=instance) - - tracer = analysis.tracer_via_instance_from(instance=instance) - - fit = al.FitInterferometer(dataset=interferometer_7, tracer=tracer) - - assert fit.log_likelihood == analysis_log_likelihood - assert analysis_log_likelihood == pytest.approx(-62.463179940, 1.0e-4) - - positions_likelihood = al.PositionsLH( - positions=al.Grid2DIrregular([(1.0, 100.0), (200.0, 2.0)]), threshold=0.01 - ) - - analysis = al.AnalysisInterferometer( - dataset=interferometer_7, - positions_likelihood_list=[positions_likelihood], - use_jax=False, - ) - analysis_log_likelihood = analysis.log_likelihood_function(instance=instance) - - assert analysis_log_likelihood == pytest.approx(-44097289569.2342, 1.0e-4) - - -def _pixelization_model(): - pixelization = al.Pixelization( - mesh=al.mesh.RectangularUniform(shape=(3, 3)), - regularization=al.reg.Constant(coefficient=0.01), - ) - return af.Collection( - galaxies=af.Collection( - galaxy_0=al.Galaxy(redshift=0.5), - galaxy_1=al.Galaxy(redshift=0.5, pixelization=pixelization), - ) - ) - - -def test__shared_state_from__preloads_curvature_reused__figure_of_merit_unchanged( - interferometer_7, -): - dataset = interferometer_7.apply_sparse_operator(use_jax=False) - - model = _pixelization_model() - instance = model.instance_from_unit_vector([]) - - analysis = al.AnalysisInterferometer( - dataset=dataset, use_jax=False, shared_preloads=True - ) - - # `shared_state_from` builds a `PreloadsInterferometer` carrying the curvature matrix `F` and - # the mapper (the channel-invariant inversion-setup quantities). - shared = analysis.shared_state_from(instance=instance) - assert isinstance(shared, aa.PreloadsInterferometer) - assert shared.curvature_matrix is not None - assert shared.mapper_galaxy_dict is not None - - # The preloaded `F` and mapper are reused by the fit (identity) and leave the figure of merit - # unchanged. Reusing the mapper means the Delaunay triangulation is not rebuilt per channel. - fit_unshared = analysis.fit_from(instance=instance) - fit_shared = analysis.fit_from(instance=instance, preloads=shared) - - assert fit_shared.inversion.curvature_matrix is shared.curvature_matrix - assert fit_shared.tracer_to_inversion.mapper_galaxy_dict is shared.mapper_galaxy_dict - assert fit_shared.figure_of_merit == pytest.approx(fit_unshared.figure_of_merit) - - # The full `log_likelihood_function` with the shared object matches the unshared call. - assert analysis.log_likelihood_function( - instance=instance, shared=shared - ) == pytest.approx(analysis.log_likelihood_function(instance=instance)) - - -def test__shared_state_from__returns_none_when_not_opted_in(interferometer_7): - dataset = interferometer_7.apply_sparse_operator(use_jax=False) - - model = _pixelization_model() - instance = model.instance_from_unit_vector([]) - - analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=False) - - assert analysis.shared_state_from(instance=instance) is None - - -def test__preloads_scoped__cross_type_preloads_reduced_to_mesh_view(interferometer_7): - lens = al.Galaxy(redshift=0.5, light=al.lp.Sersic(intensity=0.1)) - tracer = al.Tracer(galaxies=[lens]) - - # Cross-dataset-type preloads (e.g. from an imaging lead factor in a joint graph): only - # the mesh-geometry view is valid for an interferometer fit. - cross_type = aa.PreloadsImaging( - source_plane_mesh_grid=[["mesh"]], image_plane_mesh_grid=[["image-mesh"]] - ) - - fit = al.FitInterferometer( - dataset=interferometer_7, tracer=tracer, preloads=cross_type - ) - - scoped = fit._preloads_scoped - assert isinstance(scoped, aa.PreloadsInterferometer) - assert scoped.source_plane_mesh_grid == [["mesh"]] - assert scoped.curvature_matrix is None - assert scoped.mapper_galaxy_dict is None - - same_type = aa.PreloadsInterferometer(curvature_matrix="F") - fit = al.FitInterferometer( - dataset=interferometer_7, tracer=tracer, preloads=same_type - ) - assert fit._preloads_scoped is same_type - - -def test__shared_state_from__populates_mesh_geometry_fields(interferometer_7): - dataset = interferometer_7.apply_sparse_operator(use_jax=False) - - model = _pixelization_model() - instance = model.instance_from_unit_vector([]) - - analysis = al.AnalysisInterferometer( - dataset=dataset, use_jax=False, shared_preloads=True - ) - - # The mesh-geometry fields ride alongside the curvature matrix + mapper so that - # cross-dataset-type factors of a joint graph can consume the shared mesh. - shared = analysis.shared_state_from(instance=instance) - assert shared.source_plane_mesh_grid is not None - assert shared.image_plane_mesh_grid is not None +import importlib.util +from pathlib import Path +import pytest + +import autofit as af +import autoarray as aa +import autolens as al +from autolens import exc + +# The interferometer shared-state preload builds the jax-backed NUFFT sparse +# operator, so these tests need jax installed to run (it ships via the +# `[optional]` extras). The NumPy-only Python-version matrix has no jax, so skip +# there rather than fail. +requires_jax = pytest.mark.skipif( + importlib.util.find_spec("jax") is None, + reason="requires jax (installed via the [optional] extras; absent on the NumPy-only matrix env)", +) + +from autolens.interferometer.model.result import ResultInterferometer + +directory = Path(__file__).resolve().parent + + +def test__make_result__result_interferometer_is_returned(interferometer_7): + model = af.Collection(galaxies=af.Collection(galaxy_0=al.Galaxy(redshift=0.5))) + + analysis = al.AnalysisInterferometer(dataset=interferometer_7, use_jax=False) + + search = al.m.MockSearch(name="test_search") + + result = search.fit(model=model, analysis=analysis) + + assert isinstance(result, ResultInterferometer) + + +def test__figure_of_merit__matches_correct_fit_given_galaxy_profiles(interferometer_7): + lens_galaxy = al.Galaxy(redshift=0.5, light=al.lp.Sersic(intensity=0.1)) + + model = af.Collection(galaxies=af.Collection(lens=lens_galaxy)) + + analysis = al.AnalysisInterferometer(dataset=interferometer_7, use_jax=False) + + instance = model.instance_from_unit_vector([]) + analysis_log_likelihood = analysis.log_likelihood_function(instance=instance) + + tracer = analysis.tracer_via_instance_from(instance=instance) + + fit = al.FitInterferometer(dataset=interferometer_7, tracer=tracer) + + assert fit.log_likelihood == analysis_log_likelihood + + +def test__positions__likelihood_overwrite__changes_likelihood( + interferometer_7, mask_2d_7x7 +): + lens = al.Galaxy(redshift=0.5, mass=al.mp.IsothermalSph(centre=(0.05, 0.05))) + source = al.Galaxy(redshift=1.0, light=al.lp.SersicSph(centre=(0.05, 0.05))) + + model = af.Collection(galaxies=af.Collection(lens=lens, source=source)) + + analysis = al.AnalysisInterferometer(dataset=interferometer_7, use_jax=False) + + instance = model.instance_from_unit_vector([]) + analysis_log_likelihood = analysis.log_likelihood_function(instance=instance) + + tracer = analysis.tracer_via_instance_from(instance=instance) + + fit = al.FitInterferometer(dataset=interferometer_7, tracer=tracer) + + assert fit.log_likelihood == analysis_log_likelihood + assert analysis_log_likelihood == pytest.approx(-62.463179940, 1.0e-4) + + positions_likelihood = al.PositionsLH( + positions=al.Grid2DIrregular([(1.0, 100.0), (200.0, 2.0)]), threshold=0.01 + ) + + analysis = al.AnalysisInterferometer( + dataset=interferometer_7, + positions_likelihood_list=[positions_likelihood], + use_jax=False, + ) + analysis_log_likelihood = analysis.log_likelihood_function(instance=instance) + + assert analysis_log_likelihood == pytest.approx(-44097289569.2342, 1.0e-4) + + +def _pixelization_model(): + pixelization = al.Pixelization( + mesh=al.mesh.RectangularUniform(shape=(3, 3)), + regularization=al.reg.Constant(coefficient=0.01), + ) + return af.Collection( + galaxies=af.Collection( + galaxy_0=al.Galaxy(redshift=0.5), + galaxy_1=al.Galaxy(redshift=0.5, pixelization=pixelization), + ) + ) + + +@requires_jax +def test__shared_state_from__preloads_curvature_reused__figure_of_merit_unchanged( + interferometer_7, +): + dataset = interferometer_7.apply_sparse_operator(use_jax=False) + + model = _pixelization_model() + instance = model.instance_from_unit_vector([]) + + analysis = al.AnalysisInterferometer( + dataset=dataset, use_jax=False, shared_preloads=True + ) + + # `shared_state_from` builds a `PreloadsInterferometer` carrying the curvature matrix `F` and + # the mapper (the channel-invariant inversion-setup quantities). + shared = analysis.shared_state_from(instance=instance) + assert isinstance(shared, aa.PreloadsInterferometer) + assert shared.curvature_matrix is not None + assert shared.mapper_galaxy_dict is not None + + # The preloaded `F` and mapper are reused by the fit (identity) and leave the figure of merit + # unchanged. Reusing the mapper means the Delaunay triangulation is not rebuilt per channel. + fit_unshared = analysis.fit_from(instance=instance) + fit_shared = analysis.fit_from(instance=instance, preloads=shared) + + assert fit_shared.inversion.curvature_matrix is shared.curvature_matrix + assert fit_shared.tracer_to_inversion.mapper_galaxy_dict is shared.mapper_galaxy_dict + assert fit_shared.figure_of_merit == pytest.approx(fit_unshared.figure_of_merit) + + # The full `log_likelihood_function` with the shared object matches the unshared call. + assert analysis.log_likelihood_function( + instance=instance, shared=shared + ) == pytest.approx(analysis.log_likelihood_function(instance=instance)) + + +@requires_jax +def test__shared_state_from__returns_none_when_not_opted_in(interferometer_7): + dataset = interferometer_7.apply_sparse_operator(use_jax=False) + + model = _pixelization_model() + instance = model.instance_from_unit_vector([]) + + analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=False) + + assert analysis.shared_state_from(instance=instance) is None + + +def test__preloads_scoped__cross_type_preloads_reduced_to_mesh_view(interferometer_7): + lens = al.Galaxy(redshift=0.5, light=al.lp.Sersic(intensity=0.1)) + tracer = al.Tracer(galaxies=[lens]) + + # Cross-dataset-type preloads (e.g. from an imaging lead factor in a joint graph): only + # the mesh-geometry view is valid for an interferometer fit. + cross_type = aa.PreloadsImaging( + source_plane_mesh_grid=[["mesh"]], image_plane_mesh_grid=[["image-mesh"]] + ) + + fit = al.FitInterferometer( + dataset=interferometer_7, tracer=tracer, preloads=cross_type + ) + + scoped = fit._preloads_scoped + assert isinstance(scoped, aa.PreloadsInterferometer) + assert scoped.source_plane_mesh_grid == [["mesh"]] + assert scoped.curvature_matrix is None + assert scoped.mapper_galaxy_dict is None + + same_type = aa.PreloadsInterferometer(curvature_matrix="F") + fit = al.FitInterferometer( + dataset=interferometer_7, tracer=tracer, preloads=same_type + ) + assert fit._preloads_scoped is same_type + + +@requires_jax +def test__shared_state_from__populates_mesh_geometry_fields(interferometer_7): + dataset = interferometer_7.apply_sparse_operator(use_jax=False) + + model = _pixelization_model() + instance = model.instance_from_unit_vector([]) + + analysis = al.AnalysisInterferometer( + dataset=dataset, use_jax=False, shared_preloads=True + ) + + # The mesh-geometry fields ride alongside the curvature matrix + mapper so that + # cross-dataset-type factors of a joint graph can consume the shared mesh. + shared = analysis.shared_state_from(instance=instance) + assert shared.source_plane_mesh_grid is not None + assert shared.image_plane_mesh_grid is not None diff --git a/test_autolens/lens/test_substructure_util.py b/test_autolens/lens/test_substructure_util.py index da76b80b0..95d969cb5 100644 --- a/test_autolens/lens/test_substructure_util.py +++ b/test_autolens/lens/test_substructure_util.py @@ -1,3 +1,5 @@ +import importlib.util + import numpy as np import pytest @@ -10,6 +12,14 @@ reason="Kaplinghat SIDM profiles are provided by the pending PyAutoGalaxy release.", ) +# `galaxies_to_halo_arrays` is jax-backed; these tests need jax installed to run +# (it ships via the `[optional]` extras). The NumPy-only Python-version matrix +# has no jax, so skip there rather than fail. +requires_jax = pytest.mark.skipif( + importlib.util.find_spec("jax") is None, + reason="requires jax (installed via the [optional] extras; absent on the NumPy-only matrix env)", +) + @requires_kaplinghat def test__autolens_exposes_kaplinghat_profiles_from_autogalaxy(): @@ -17,6 +27,7 @@ def test__autolens_exposes_kaplinghat_profiles_from_autogalaxy(): assert hasattr(al.mp, "KaplinghatCoredNFWMCRLudlowSph") +@requires_jax @requires_kaplinghat def test__galaxies_to_halo_arrays__packs_kaplinghat_profile_parameters(): profile = al.mp.KaplinghatCoredNFWSph( @@ -59,6 +70,7 @@ def test__galaxies_to_halo_arrays__packs_kaplinghat_profile_parameters(): ) +@requires_jax def test__galaxies_to_halo_arrays__raises_for_unsupported_profile_class(): with pytest.raises(ValueError): substructure_util.galaxies_to_halo_arrays(