From 2de8c7b16b09baff7229a29aa2d437636e579025 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 30 Jul 2026 19:59:19 +0100 Subject: [PATCH] =?UTF-8?q?feat:=20solved=20point-source=20fit=20coverage?= =?UTF-8?q?=20=E2=80=94=20jax=5Flikelihood=20variants,=20gradients,=20flux?= =?UTF-8?q?es+time-delays=20(#657=20phase=203)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...t_dataset_with_fluxes_and_time_delays.json | 98 +++++++ scripts/CLAUDE.md | 7 +- scripts/point_source/jax_grad/gradient.py | 83 ++++++ .../jax_likelihood/fluxes_time_delays.py | 270 ++++++++++++++++++ .../jax_likelihood/image_plane.py | 186 ++++++++++++ scripts/point_source/jax_likelihood/point.py | 129 +++++++-- .../jax_likelihood/source_plane.py | 191 +++++++++++-- scripts/point_source/simulators/simple.py | 52 +++- 8 files changed, 970 insertions(+), 46 deletions(-) create mode 100644 dataset/point_source/simple/point_dataset_with_fluxes_and_time_delays.json create mode 100644 scripts/point_source/jax_likelihood/fluxes_time_delays.py diff --git a/dataset/point_source/simple/point_dataset_with_fluxes_and_time_delays.json b/dataset/point_source/simple/point_dataset_with_fluxes_and_time_delays.json new file mode 100644 index 00000000..084b523e --- /dev/null +++ b/dataset/point_source/simple/point_dataset_with_fluxes_and_time_delays.json @@ -0,0 +1,98 @@ +{ + "type": "instance", + "class_path": "autolens.point.dataset.PointDataset", + "arguments": { + "name": "point_0", + "positions": { + "type": "instance", + "class_path": "autoarray.structures.grids.irregular_2d.Grid2DIrregular", + "arguments": { + "values": { + "type": "ndarray", + "array": [ + [ + 1.3726562500000001, + 0.9936739398630823 + ], + [ + -0.9460937500000001, + -1.1578579226638823 + ] + ], + "dtype": "float64" + } + } + }, + "fluxes": { + "type": "instance", + "class_path": "autoarray.structures.arrays.irregular.ArrayIrregular", + "arguments": { + "values": { + "type": "ndarray", + "array": [ + 23.7215548997517, + 13.79784180512648 + ], + "dtype": "float64" + } + } + }, + "redshift": null, + "time_delays": { + "type": "instance", + "class_path": "autoarray.structures.arrays.irregular.ArrayIrregular", + "arguments": { + "values": { + "type": "ndarray", + "array": [ + -179.6711648183478, + -144.82858583437982 + ], + "dtype": "float64" + } + } + }, + "positions_noise_map": { + "type": "instance", + "class_path": "autoarray.structures.arrays.irregular.ArrayIrregular", + "arguments": { + "values": { + "type": "ndarray", + "array": [ + 0.005, + 0.005 + ], + "dtype": "float64" + } + } + }, + "fluxes_noise_map": { + "type": "instance", + "class_path": "autoarray.structures.arrays.irregular.ArrayIrregular", + "arguments": { + "values": { + "type": "ndarray", + "array": [ + 2.37215548997517, + 1.379784180512648 + ], + "dtype": "float64" + } + } + }, + "time_delays_noise_map": { + "type": "instance", + "class_path": "autoarray.structures.arrays.irregular.ArrayIrregular", + "arguments": { + "values": { + "type": "ndarray", + "array": [ + 0.5, + 0.5 + ], + "dtype": "float64" + } + } + } + } +} \ No newline at end of file diff --git a/scripts/CLAUDE.md b/scripts/CLAUDE.md index b87e54ee..d8d37071 100644 --- a/scripts/CLAUDE.md +++ b/scripts/CLAUDE.md @@ -165,9 +165,10 @@ likelihoods of the same base name. | `interferometer/jax_likelihood/rectangular_mge.py` | Rectangular source + MGE lens for interferometry | | `interferometer/jax_likelihood/rectangular_dspl.py` | Rectangular source on double source plane (interferometry) | | `interferometer/jax_likelihood/rectangular_sparse.py` | Rectangular pixelization via JAX sparse-operator NUFFT path | -| `point_source/jax_likelihood/point.py` | Point-source likelihood | -| `point_source/jax_likelihood/image_plane.py` | Point-source image-plane chi-squared (`FitPositionsImagePairAll`) | -| `point_source/jax_likelihood/source_plane.py` | Point-source source-plane chi-squared (`FitPositionsSource`) — JIT currently blocked | +| `point_source/jax_likelihood/point.py` | Point-source likelihood walkthrough (image-plane `FitPositionsImagePairAll`) + centre-free `FitPositionsImagePairAllSolved` coverage | +| `point_source/jax_likelihood/image_plane.py` | Point-source image-plane chi-squared (`FitPositionsImagePairAll`) + centre-free `FitPositionsImagePairAllSolved` / `FitPositionsImagePairRepeatSolved` variants | +| `point_source/jax_likelihood/source_plane.py` | Point-source source-plane chi-squared (`FitPositionsSource`) + centre-free `FitPositionsSourceSolved` — Path A JIT blocked by the fit-return pytree gap (`PyAutoPrompt/autolens/fit_point_pytree.md`), not the (already-fixed) xp-propagation bug | +| `point_source/jax_likelihood/fluxes_time_delays.py` | Point-source fluxes + time delays via the solved fit classes (`FitFluxesSolved`, `FitTimeDelaysSolved`) alongside `FitPositionsSourceSolved` | | `multi/jax_likelihood/lp.py` | Parametric Sersic across g/r via `FactorGraphModel`; per-band source `ell_comps` (option B) | | `multi/jax_likelihood/mge.py` | MGE source across g/r; per-band source MGE `ell_comps` (option B) | | `multi/jax_likelihood/mge_group.py` | MGE + extra galaxies across g/r; per-band source MGE `ell_comps` (option B) | diff --git a/scripts/point_source/jax_grad/gradient.py b/scripts/point_source/jax_grad/gradient.py index 11bb620c..4df0d4a0 100644 --- a/scripts/point_source/jax_grad/gradient.py +++ b/scripts/point_source/jax_grad/gradient.py @@ -183,3 +183,86 @@ ), "A positional parameter has zero gradient — evaluation point is degenerate." print("point_source.py JAX gradient checks passed.") + + +""" +__Solved Source-Plane Gradient (Parameter-Free Centre)__ + +Repeats the finiteness + finite-difference checks above for +``al.FitPositionsSourceSolved`` against a parameter-free ``al.ps.PointSolved`` +source. With no source-centre parameters, the model's only positional degrees +of freedom are the lens mass parameters; ``cosmology.H0`` still has no +bearing on this position-only chi-squared and stays possibly-zero-grad. + +Note: gradients through the image-plane ``PointSolver`` variants +(``FitPositionsImagePairAll`` / ``...Solved`` etc.) require a ``custom_jvp`` +around the triangle-refinement solve and are deliberately not attempted here +— that is phase 5 of issue #657. +""" + +point_0_solved = af.Model(al.ps.PointSolved) + +source_solved = af.Model(al.Galaxy, redshift=1.0, point_0=point_0_solved) + +model_solved = af.Collection( + galaxies=af.Collection(lens=lens, source=source_solved), cosmology=cosmology +) + +print(model_solved.info) + +analysis_solved = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsSourceSolved, +) + +fitness_solved = Fitness( + model=model_solved, + analysis=analysis_solved, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99, +) + +param_vector_solved = jnp.array(model_solved.physical_values_from_prior_medians) + +key_solved = jax.random.PRNGKey(43) +perturbation_solved = jax.random.uniform( + key_solved, shape=param_vector_solved.shape, minval=0.001, maxval=0.005 +) +param_vector_solved = param_vector_solved + perturbation_solved + +value_solved, grad_solved = jax.value_and_grad(fitness_solved.call)(param_vector_solved) + +print(f"Log likelihood (solved) = {float(value_solved):.6f}") +print(f"Gradient shape (solved) = {grad_solved.shape}") + +assert np.isfinite(float(value_solved)), "Log likelihood (solved) is not finite" +assert grad_solved.shape == ( + model_solved.total_free_parameters, +), f"Gradient shape mismatch (solved): {grad_solved.shape}" +assert np.all( + np.isfinite(np.array(grad_solved)) +), f"Gradient contains non-finite values (solved): {np.array(grad_solved)}" +assert not np.all(np.array(grad_solved) == 0.0), "Gradient is all zeros (solved)" + +param_names_solved = util.parameter_names_from(model_solved) + +comparison_solved = util.compare_gradients( + fitness_solved.call, + param_vector_solved, + param_names=param_names_solved, +) + +util.assert_gradients_match(comparison_solved) + +# With PointSolved there are no source-centre parameters — the only +# positional degrees of freedom are the lens mass parameters. H0 must still +# be excluded (this position-only chi-squared has no dependence on it). +positional_indices_solved = [ + i for i, name in enumerate(param_names_solved) if "H0" not in name +] +assert np.all( + np.abs(comparison_solved["ad"][positional_indices_solved]) > 0.0 +), "A positional parameter has zero gradient — evaluation point is degenerate (solved)." + +print("point_source gradient.py solved-source checks passed.") diff --git a/scripts/point_source/jax_likelihood/fluxes_time_delays.py b/scripts/point_source/jax_likelihood/fluxes_time_delays.py new file mode 100644 index 00000000..c20b0695 --- /dev/null +++ b/scripts/point_source/jax_likelihood/fluxes_time_delays.py @@ -0,0 +1,270 @@ +""" +Func Grad: Point Source Fluxes + Time Delays (Solved) +======================================================== + +Test that JAX can compute the log-likelihood of a ``PointDataset`` that +carries fluxes and time delays in addition to positions, using the solved +(parameter-free) point-source fit classes throughout: +``al.FitPositionsSourceSolved`` for the source-plane position chi-squared, +``al.FitFluxesSolved`` for the flux chi-squared, and ``al.FitTimeDelaysSolved`` +for the time-delay chi-squared — all evaluated against the same +analytically-solved source centre β* (Lombardi 2024, arXiv:2406.15280). + +Because ``al.ps.PointSolved`` is parameter-free, the vmap model composes to +just the lens mass parameters plus a free ``cosmology.H0`` — H0 is kept free +because time delays are the observable that constrains it. + +Full-pipeline JIT status +------------------------ + +Path A (``jax.jit(analysis.fit_from)``) goes through the same source-plane +fit as ``point_source/jax_likelihood/source_plane.py`` and is gated by the +same fit-return pytree gap: ``fit_from`` returns a bare ``PointSolver`` at +output component ``[1][1]``, tracked in +``PyAutoPrompt/autolens/fit_point_pytree.md``. When Path A JIT fails with +this ``TypeError`` the script prints a clear BLOCKER line and continues. + +__Env__ + +Test-harness configuration (PyAutoHands docs/env_profile_redesign.md §10). +JAX likelihood functions test JIT compilation; need JAX enabled and full- +size datasets. + +ENV: jax full_datasets +""" + +import numpy as np +import jax.numpy as jnp +import jax +from pathlib import Path + +import autofit as af +import autolens as al + + +""" +__Dataset__ +""" +dataset_name = "simple" +dataset_path = Path("dataset") / "point_source" / dataset_name +dataset_file = dataset_path / "point_dataset_with_fluxes_and_time_delays.json" + +""" +__Dataset Auto-Simulation__ + +Guarded on this file specifically (not ``al.util.dataset.should_simulate``, +which only checks the directory): the directory already exists once +``simple.py``'s positions-only output is present, but the flux/time-delay +file is a second output of the same simulator added alongside it. +""" +if not dataset_file.exists(): + import subprocess + import sys + + subprocess.run( + [sys.executable, "scripts/point_source/simulators/simple.py"], + check=True, + ) + +dataset = al.from_json(file_path=dataset_file) + +""" +__Point Solver__ +""" +grid = al.Grid2D.uniform( + shape_native=(100, 100), + pixel_scales=0.2, +) + +solver = al.PointSolver.for_grid( + grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1 +) + +""" +__Model__ + +Mass priors copied from ``point_source/jax_likelihood/source_plane.py``. The +source is parameter-free (``al.ps.PointSolved``); cosmology stays free +(``H0``) because the time delays constrain it. +""" +mass = af.Model(al.mp.Isothermal) + +mass.centre.centre_0 = af.UniformPrior(lower_limit=0.0, upper_limit=0.02) +mass.centre.centre_1 = af.UniformPrior(lower_limit=0.0, upper_limit=0.02) +mass.ell_comps.ell_comps_0 = af.UniformPrior(lower_limit=0.0, upper_limit=0.02) +mass.ell_comps.ell_comps_1 = af.UniformPrior(lower_limit=0.0, upper_limit=0.02) +mass.einstein_radius = af.UniformPrior(lower_limit=1.5, upper_limit=1.8) + +lens = af.Model(al.Galaxy, redshift=0.5, mass=mass) + +point_0 = af.Model(al.ps.PointSolved) + +source = af.Model(al.Galaxy, redshift=1.0, point_0=point_0) + +cosmology = af.Model(al.cosmo.FlatLambdaCDM) +cosmology.H0 = af.UniformPrior(lower_limit=0.0, upper_limit=150.0) + +model = af.Collection( + galaxies=af.Collection(lens=lens, source=source), cosmology=cosmology +) + +print(model.info) + +""" +__Analysis__ +""" +analysis = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsSourceSolved, + fit_flux_cls=al.FitFluxesSolved, + fit_time_delays_cls=al.FitTimeDelaysSolved, +) + +from autofit.non_linear.fitness import Fitness +import time + +batch_size = 1 + +fitness = Fitness( + model=model, + analysis=analysis, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99, +) + +parameters = np.zeros((batch_size, model.total_free_parameters)) +for i in range(batch_size): + parameters[i, :] = model.physical_values_from_prior_medians +parameters = jnp.array(parameters) + +start = time.time() +print() +print(fitness._vmap(parameters)) +print("JAX Time To VMAP + JIT Function", time.time() - start) + +start = time.time() +print() +result = fitness._vmap(parameters) +print(result) +print("JAX Time Taken using VMAP:", time.time() - start) +print("JAX Time Taken per Likelihood:", (time.time() - start) / batch_size) + +EXPECTED_VMAP_LOG_LIKELIHOOD_FLUXES_TIME_DELAYS = -137.67455417 + +np.testing.assert_allclose( + np.array(result), + EXPECTED_VMAP_LOG_LIKELIHOOD_FLUXES_TIME_DELAYS, + rtol=1e-4, + err_msg="point_source/fluxes_time_delays: JAX vmap likelihood mismatch", +) + + +""" +__Path A: jit-wrap ``analysis.fit_from``__ + +Free ``cosmology`` is dropped for this block (same ``UnexpectedTracerError`` +caveat as ``image_plane.py`` / ``point.py``) even though ``H0`` is free in +the vmap model above. +""" + +model_jit = af.Collection(galaxies=af.Collection(lens=lens, source=source)) + +instance = model_jit.instance_from_prior_medians() + +analysis_np = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsSourceSolved, + fit_flux_cls=al.FitFluxesSolved, + fit_time_delays_cls=al.FitTimeDelaysSolved, + use_jax=False, +) +fit_np = analysis_np.fit_from(instance=instance) +log_likelihood_np = float(fit_np.log_likelihood) +print("NumPy fit.log_likelihood:", log_likelihood_np) + +# This is the Path-A *cosmology-dropped* reference value (model_jit has no +# cosmology component, so AnalysisPoint falls back to its internal default +# rather than the vmap model's free H0=75 prior median) — it is NOT expected +# to match the vmap literal above; see the H0-matched parity check below. +EXPECTED_EAGER_LOG_LIKELIHOOD_FLUXES_TIME_DELAYS = -133.0115150576916 + +np.testing.assert_allclose( + log_likelihood_np, + EXPECTED_EAGER_LOG_LIKELIHOOD_FLUXES_TIME_DELAYS, + rtol=1e-4, + err_msg=( + f"point_source/fluxes_time_delays: regression — eager (cosmology-dropped, " + f"Path-A reference) log_likelihood drifted (got {log_likelihood_np}, " + f"expected {EXPECTED_EAGER_LOG_LIKELIHOOD_FLUXES_TIME_DELAYS})" + ), +) +print( + f"Eager regression assertion PASSED (cosmology-dropped Path-A reference): " + f"log_likelihood matches {EXPECTED_EAGER_LOG_LIKELIHOOD_FLUXES_TIME_DELAYS}" +) + +""" +__NumPy-vs-JAX vmap parity (H0-matched)__ + +Time delays are directly and strongly H0-sensitive, so eager/vmap parity is +only meaningful when both paths share the same cosmology — unlike the +positions-only scripts (e.g. ``source_plane.py``), where dropping cosmology +in Path A is a no-op on the likelihood and parity holds incidentally. Here we +build a second eager fit from the *full* vmap model's prior-median instance +(cosmology included, H0=75) rather than reusing the cosmology-dropped +``instance`` above. +""" +instance_full = model.instance_from_prior_medians() + +analysis_full_np = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsSourceSolved, + fit_flux_cls=al.FitFluxesSolved, + fit_time_delays_cls=al.FitTimeDelaysSolved, + use_jax=False, +) +fit_full_np = analysis_full_np.fit_from(instance=instance_full) +log_likelihood_full_np = float(fit_full_np.log_likelihood) +print("NumPy fit.log_likelihood (H0-matched, full model):", log_likelihood_full_np) + +np.testing.assert_allclose( + log_likelihood_full_np, + float(result[0]), + rtol=1e-4, + err_msg="point_source/fluxes_time_delays: eager vs vmap parity mismatch (H0-matched)", +) +print("PASS: eager (H0-matched) vs vmap parity.") + +analysis_jit = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsSourceSolved, + fit_flux_cls=al.FitFluxesSolved, + fit_time_delays_cls=al.FitTimeDelaysSolved, + use_jax=True, +) +fit_jit_fn = jax.jit(analysis_jit.fit_from) + +full_pipeline_jits = False +try: + fit = fit_jit_fn(instance) + print("JIT fit.log_likelihood:", fit.log_likelihood) + assert isinstance( + fit.log_likelihood, jnp.ndarray + ), f"expected jax.Array, got {type(fit.log_likelihood)}" + np.testing.assert_allclose(float(fit.log_likelihood), log_likelihood_np, rtol=1e-4) + full_pipeline_jits = True + print("PASS: jit(fit_from) round-trip matches NumPy scalar.") +except TypeError as e: + print( + "\nBLOCKER: fluxes/time-delays jit(fit_from) is gated by:\n" + f" {type(e).__name__}: {e}\n" + " fit_from returns a bare PointSolver at output component [1][1],\n" + " which is not pytree-registered under jax.jit (it goes through the\n" + " same source-plane fit as source_plane.py). Tracked in\n" + " PyAutoPrompt/autolens/fit_point_pytree.md.\n" + " Eager NumPy regression assertion still PASSED above." + ) diff --git a/scripts/point_source/jax_likelihood/image_plane.py b/scripts/point_source/jax_likelihood/image_plane.py index fb09d009..d62e229c 100644 --- a/scripts/point_source/jax_likelihood/image_plane.py +++ b/scripts/point_source/jax_likelihood/image_plane.py @@ -185,3 +185,189 @@ float(fit.log_likelihood), float(fit_np.log_likelihood), rtol=1e-4 ) print("PASS: jit(fit_from) round-trip matches NumPy scalar.") + + +""" +__Model: Solved Source (Parameter-Free)__ + +Swaps the source for parameter-free ``al.ps.PointSolved`` to exercise the +centre-free image-plane variants below. +""" + +point_0_solved = af.Model(al.ps.PointSolved) + +source_solved = af.Model(al.Galaxy, redshift=1.0, point_0=point_0_solved) + +model_solved = af.Collection( + galaxies=af.Collection(lens=lens, source=source_solved), cosmology=cosmology +) + +print(model_solved.info) + +model_solved_jit = af.Collection( + galaxies=af.Collection(lens=lens, source=source_solved) +) +instance_solved = model_solved_jit.instance_from_prior_medians() + +""" +__Analysis: FitPositionsImagePairAllSolved__ +""" +analysis_all_solved = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsImagePairAllSolved, +) + +from autofit.non_linear.fitness import Fitness +import time + +batch_size = 1 + +fitness_all_solved = Fitness( + model=model_solved, + analysis=analysis_all_solved, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99, +) + +parameters_all_solved = np.zeros((batch_size, model_solved.total_free_parameters)) +for i in range(batch_size): + parameters_all_solved[i, :] = model_solved.physical_values_from_prior_medians +parameters_all_solved = jnp.array(parameters_all_solved) + +start = time.time() +print() +print(fitness_all_solved._vmap(parameters_all_solved)) +print("JAX Time To VMAP + JIT Function", time.time() - start) + +start = time.time() +print() +result_all_solved = fitness_all_solved._vmap(parameters_all_solved) +print(result_all_solved) +print("JAX Time Taken using VMAP:", time.time() - start) +print("JAX Time Taken per Likelihood:", (time.time() - start) / batch_size) + +EXPECTED_VMAP_LOG_LIKELIHOOD_IMAGE_PLANE_ALL_SOLVED = -82.33883111 + +np.testing.assert_allclose( + np.array(result_all_solved), + EXPECTED_VMAP_LOG_LIKELIHOOD_IMAGE_PLANE_ALL_SOLVED, + rtol=1e-4, + err_msg="point_source/image_plane: JAX vmap likelihood mismatch (all solved)", +) + + +""" +__Path A: jit-wrap ``analysis.fit_from`` (FitPositionsImagePairAllSolved)__ + +Expected to JIT end-to-end exactly like the modelled-centre ``FitPositionsImagePairAll`` +block above — no try/except; a failure here is a real regression. +""" +analysis_all_solved_np = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsImagePairAllSolved, + use_jax=False, +) +fit_all_solved_np = analysis_all_solved_np.fit_from(instance=instance_solved) +print("NumPy fit.log_likelihood (all solved):", float(fit_all_solved_np.log_likelihood)) + +analysis_all_solved_jit = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsImagePairAllSolved, + use_jax=True, +) +fit_all_solved_jit_fn = jax.jit(analysis_all_solved_jit.fit_from) +fit_all_solved = fit_all_solved_jit_fn(instance_solved) + +print("JIT fit.log_likelihood (all solved):", fit_all_solved.log_likelihood) +assert isinstance( + fit_all_solved.log_likelihood, jnp.ndarray +), f"expected jax.Array, got {type(fit_all_solved.log_likelihood)}" +np.testing.assert_allclose( + float(fit_all_solved.log_likelihood), float(fit_all_solved_np.log_likelihood), rtol=1e-4 +) +print("PASS: jit(fit_from) round-trip matches NumPy scalar (all solved).") + + +""" +__Analysis: FitPositionsImagePairRepeatSolved__ +""" +analysis_repeat_solved = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsImagePairRepeatSolved, +) + +fitness_repeat_solved = Fitness( + model=model_solved, + analysis=analysis_repeat_solved, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99, +) + +parameters_repeat_solved = np.zeros((batch_size, model_solved.total_free_parameters)) +for i in range(batch_size): + parameters_repeat_solved[i, :] = model_solved.physical_values_from_prior_medians +parameters_repeat_solved = jnp.array(parameters_repeat_solved) + +start = time.time() +print() +print(fitness_repeat_solved._vmap(parameters_repeat_solved)) +print("JAX Time To VMAP + JIT Function", time.time() - start) + +start = time.time() +print() +result_repeat_solved = fitness_repeat_solved._vmap(parameters_repeat_solved) +print(result_repeat_solved) +print("JAX Time Taken using VMAP:", time.time() - start) +print("JAX Time Taken per Likelihood:", (time.time() - start) / batch_size) + +EXPECTED_VMAP_LOG_LIKELIHOOD_IMAGE_PLANE_REPEAT_SOLVED = -89.71129442 + +np.testing.assert_allclose( + np.array(result_repeat_solved), + EXPECTED_VMAP_LOG_LIKELIHOOD_IMAGE_PLANE_REPEAT_SOLVED, + rtol=1e-4, + err_msg="point_source/image_plane: JAX vmap likelihood mismatch (repeat solved)", +) + + +""" +__Path A: jit-wrap ``analysis.fit_from`` (FitPositionsImagePairRepeatSolved)__ + +Expected to JIT end-to-end exactly like the modelled-centre ``FitPositionsImagePairAll`` +block above — no try/except; a failure here is a real regression. +""" +analysis_repeat_solved_np = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsImagePairRepeatSolved, + use_jax=False, +) +fit_repeat_solved_np = analysis_repeat_solved_np.fit_from(instance=instance_solved) +print( + "NumPy fit.log_likelihood (repeat solved):", + float(fit_repeat_solved_np.log_likelihood), +) + +analysis_repeat_solved_jit = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsImagePairRepeatSolved, + use_jax=True, +) +fit_repeat_solved_jit_fn = jax.jit(analysis_repeat_solved_jit.fit_from) +fit_repeat_solved = fit_repeat_solved_jit_fn(instance_solved) + +print("JIT fit.log_likelihood (repeat solved):", fit_repeat_solved.log_likelihood) +assert isinstance( + fit_repeat_solved.log_likelihood, jnp.ndarray +), f"expected jax.Array, got {type(fit_repeat_solved.log_likelihood)}" +np.testing.assert_allclose( + float(fit_repeat_solved.log_likelihood), + float(fit_repeat_solved_np.log_likelihood), + rtol=1e-4, +) +print("PASS: jit(fit_from) round-trip matches NumPy scalar (repeat solved).") diff --git a/scripts/point_source/jax_likelihood/point.py b/scripts/point_source/jax_likelihood/point.py index 1ffb4459..b6e8e369 100644 --- a/scripts/point_source/jax_likelihood/point.py +++ b/scripts/point_source/jax_likelihood/point.py @@ -1,27 +1,16 @@ """ -Func Grad: Light Parametric Operated -==================================== +Func Grad: Point Source Likelihood +=================================== -This script test if JAX can successfully compute the gradient of the log likelihood of an `Imaging` dataset with a -model which uses operated light profiles. +Test that JAX can compute the log-likelihood of a ``PointDataset`` using the +**image-plane** chi-squared (``al.FitPositionsImagePairAll``) via +``AnalysisPoint``, exercising both the batched ``fitness._vmap`` likelihood +and the full ``jax.jit(analysis.fit_from)`` pipeline. - __Operated Fitting__ - -It is common for galaxies to have point-source emission, for example bright emission right at their centre due to -an active galactic nuclei or very compact knot of star formation. - -This point-source emission is subject to blurring during data accquisiton due to the telescope optics, and therefore -is not seen as a single pixel of light but spread over multiple pixels as a convolution with the telescope -Point Spread Function (PSF). - -It is difficult to model this compact point source emission using a point-source light profile (or an extremely -compact Gaussian / Sersic profile). This is because when the model-image of a compact point source of light is -convolved with the PSF, the solution to this convolution is extremely sensitive to which pixel (and sub-pixel) the -compact model emission lands in. - -Operated light profiles offer an alternative approach, whereby the light profile is assumed to have already been -convolved with the PSF. This operated light profile is then fitted directly to the point-source emission, which as -discussed above shows the PSF features. +This script doubles as a documented walkthrough of the point-source modeling +API — see the ``__Point Solver__``, ``__Model__`` and ``__Name Pairing__`` +sections below — built around the same ``FitPositionsImagePairAll`` variant +tested more tersely in ``point_source/jax_likelihood/image_plane.py``. __Env__ @@ -288,3 +277,101 @@ float(fit.log_likelihood), float(fit_np.log_likelihood), rtol=1e-4 ) print("PASS: jit(fit_from) round-trip matches NumPy scalar.") + + +""" +__Model: Solved Source (Parameter-Free)__ + +Swaps the source for parameter-free ``al.ps.PointSolved`` to exercise the +centre-free ``FitPositionsImagePairAllSolved`` variant. This script is in +``smoke_tests.txt``, so only this one solved variant is added here — the +Repeat variant and the source-plane / fluxes+time-delays solved coverage live +in ``image_plane.py``, ``source_plane.py`` and ``fluxes_time_delays.py``. +""" + +point_0_solved = af.Model(al.ps.PointSolved) + +source_solved = af.Model(al.Galaxy, redshift=1.0, point_0=point_0_solved) + +model_solved = af.Collection( + galaxies=af.Collection(lens=lens, source=source_solved), cosmology=cosmology +) + +print(model_solved.info) + +analysis_all_solved = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsImagePairAllSolved, +) + +fitness_all_solved = Fitness( + model=model_solved, + analysis=analysis_all_solved, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99, +) + +parameters_all_solved = np.zeros((batch_size, model_solved.total_free_parameters)) +for i in range(batch_size): + parameters_all_solved[i, :] = model_solved.physical_values_from_prior_medians +parameters_all_solved = jnp.array(parameters_all_solved) + +start = time.time() +print() +print(fitness_all_solved._vmap(parameters_all_solved)) +print("JAX Time To VMAP + JIT Function", time.time() - start) + +start = time.time() +print() +result_all_solved = fitness_all_solved._vmap(parameters_all_solved) +print(result_all_solved) +print("JAX Time Taken using VMAP:", time.time() - start) +print("JAX Time Taken per Likelihood:", (time.time() - start) / batch_size) + +EXPECTED_VMAP_LOG_LIKELIHOOD_POINT_ALL_SOLVED = -82.33883111 + +np.testing.assert_allclose( + np.array(result_all_solved), + EXPECTED_VMAP_LOG_LIKELIHOOD_POINT_ALL_SOLVED, + rtol=1e-4, + err_msg="point: JAX vmap likelihood mismatch (all solved)", +) + + +""" +__Path A: jit-wrap ``analysis.fit_from`` (FitPositionsImagePairAllSolved)__ +""" + +model_solved_jit = af.Collection( + galaxies=af.Collection(lens=lens, source=source_solved) +) + +instance_solved = model_solved_jit.instance_from_prior_medians() + +analysis_all_solved_np = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsImagePairAllSolved, + use_jax=False, +) +fit_all_solved_np = analysis_all_solved_np.fit_from(instance=instance_solved) +print("NumPy fit.log_likelihood (all solved):", float(fit_all_solved_np.log_likelihood)) + +analysis_all_solved_jit = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsImagePairAllSolved, + use_jax=True, +) +fit_all_solved_jit_fn = jax.jit(analysis_all_solved_jit.fit_from) +fit_all_solved = fit_all_solved_jit_fn(instance_solved) + +print("JIT fit.log_likelihood (all solved):", fit_all_solved.log_likelihood) +assert isinstance( + fit_all_solved.log_likelihood, jnp.ndarray +), f"expected jax.Array, got {type(fit_all_solved.log_likelihood)}" +np.testing.assert_allclose( + float(fit_all_solved.log_likelihood), float(fit_all_solved_np.log_likelihood), rtol=1e-4 +) +print("PASS: jit(fit_from) round-trip matches NumPy scalar (all solved).") diff --git a/scripts/point_source/jax_likelihood/source_plane.py b/scripts/point_source/jax_likelihood/source_plane.py index 9c2ebba0..ec363fd7 100644 --- a/scripts/point_source/jax_likelihood/source_plane.py +++ b/scripts/point_source/jax_likelihood/source_plane.py @@ -13,14 +13,16 @@ Full-pipeline JIT status ------------------------ -The full pipeline (``jax.jit(analysis.fit_from)``) is currently **BLOCKED** -by a ``Grid2DIrregular.grid_2d_via_deflection_grid_from`` xp-propagation bug -documented in ``autolens_workspace_developer/jax_profiling/point_source/source_plane.py`` -and tracked in ``PyAutoPrompt/autolens/fit_point_pytree.md``. When Path -A JIT fails with ``TracerArrayConversionError`` the script prints a clear -BLOCKER line and continues, so the eager NumPy regression assertion is still -exercised. Once the upstream xp-propagation fix lands, the JIT path will -succeed without modifying this script. +The ``Grid2DIrregular.grid_2d_via_deflection_grid_from`` xp-propagation bug +that previously blocked Path A here was fixed in phase 2 (PyAutoArray#414). +The remaining blocker is a fit-return pytree gap: ``fit_from`` returns a +``PointSolver`` instance at output component ``[1][1]``, which is not a +valid JAX type under ``jax.jit`` — tracked in +``PyAutoPrompt/autolens/fit_point_pytree.md``. When Path A JIT fails with +this ``TypeError`` the script prints a clear BLOCKER line and continues, so +the eager NumPy regression assertion is still exercised. Once the pytree +registration lands, the JIT path will succeed without modifying this +script. __Env__ @@ -164,10 +166,11 @@ """ __Path A: jit-wrap ``analysis.fit_from``__ -Wrapped in ``try/except jax.errors.TracerArrayConversionError`` — source-plane -fitting currently fails Path A with the ``Grid2DIrregular.grid_2d_via_deflection_grid_from`` -xp-propagation bug. The eager NumPy log-likelihood is still asserted for -regression coverage. +Wrapped in ``try/except TypeError`` — source-plane fitting's ``fit_from`` +returns a bare ``PointSolver`` at output component ``[1][1]``, which is not +pytree-registered (the fit-return pytree gap tracked in +``PyAutoPrompt/autolens/fit_point_pytree.md``). The eager NumPy +log-likelihood is still asserted for regression coverage. """ @@ -219,17 +222,163 @@ np.testing.assert_allclose(float(fit.log_likelihood), log_likelihood_np, rtol=1e-4) full_pipeline_jits = True print("PASS: jit(fit_from) round-trip matches NumPy scalar.") -except (jax.errors.TracerArrayConversionError, TypeError) as e: - # Two stacked blockers gate the full-pipeline JIT: - # 1. FitPositionsSource is not pytree-registered, so fit_from returns a - # non-JAX type (TypeError). Tracked in PyAutoPrompt/autolens/fit_point_pytree.md. - # 2. Even if that is fixed, the source-plane chi-squared itself fails with - # TracerArrayConversionError owing to the Grid2DIrregular.grid_2d_via_deflection_grid_from - # xp-propagation bug (see autolens_workspace_developer/jax_profiling/point_source/source_plane.py). +except TypeError as e: + # fit_from returns a bare PointSolver instance at output component + # [1][1], which is not pytree-registered — jax.jit cannot flatten the + # return value. Tracked in PyAutoPrompt/autolens/fit_point_pytree.md. + # (The previously-blocking Grid2DIrregular.grid_2d_via_deflection_grid_from + # xp-propagation bug was fixed in phase 2, PyAutoArray#414.) print( "\nBLOCKER: source-plane jit(fit_from) is gated by:\n" f" {type(e).__name__}: {e}\n" - " Fixes tracked in PyAutoPrompt/autolens/fit_point_pytree.md and\n" - " autolens_workspace_developer/jax_profiling/point_source/source_plane.py.\n" + " fit_from returns a bare PointSolver at output component [1][1],\n" + " which is not pytree-registered under jax.jit. Tracked in\n" + " PyAutoPrompt/autolens/fit_point_pytree.md.\n" + " Eager NumPy regression assertion still PASSED above." + ) + + +""" +__Model: Solved Source (Parameter-Free)__ + +``al.ps.PointSolved`` has zero free parameters — the source-plane position β* is +solved analytically (Lombardi 2024, arXiv:2406.15280) rather than fitted, so no +centre priors are set. ``al.FitPositionsSourceSolved`` performs the same +source-plane chi-squared as ``al.FitPositionsSource`` above but against the +analytic β* instead of a modelled centre. +""" + +point_0_solved = af.Model(al.ps.PointSolved) + +source_solved = af.Model(al.Galaxy, redshift=1.0, point_0=point_0_solved) + +model_solved = af.Collection( + galaxies=af.Collection(lens=lens, source=source_solved), cosmology=cosmology +) + +print(model_solved.info) + +analysis_solved = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsSourceSolved, +) + +from autofit.non_linear.fitness import Fitness +import time + +batch_size = 1 + +fitness_solved = Fitness( + model=model_solved, + analysis=analysis_solved, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99, +) + +parameters_solved = np.zeros((batch_size, model_solved.total_free_parameters)) +for i in range(batch_size): + parameters_solved[i, :] = model_solved.physical_values_from_prior_medians +parameters_solved = jnp.array(parameters_solved) + +start = time.time() +print() +print(fitness_solved._vmap(parameters_solved)) +print("JAX Time To VMAP + JIT Function", time.time() - start) + +start = time.time() +print() +result_solved = fitness_solved._vmap(parameters_solved) +print(result_solved) +print("JAX Time Taken using VMAP:", time.time() - start) +print("JAX Time Taken per Likelihood:", (time.time() - start) / batch_size) + +EXPECTED_VMAP_LOG_LIKELIHOOD_SOURCE_PLANE_SOLVED = -94.70750993 + +np.testing.assert_allclose( + np.array(result_solved), + EXPECTED_VMAP_LOG_LIKELIHOOD_SOURCE_PLANE_SOLVED, + rtol=1e-4, + err_msg="point_source/source_plane: JAX vmap likelihood mismatch (solved)", +) + + +""" +__Path A: jit-wrap ``analysis.fit_from`` (Solved)__ + +Same narrowed ``except TypeError`` gate as the modelled-centre block above — +``fit_from`` returns a bare ``PointSolver`` at output component ``[1][1]``. +""" + +model_solved_jit = af.Collection( + galaxies=af.Collection(lens=lens, source=source_solved) +) + +instance_solved = model_solved_jit.instance_from_prior_medians() + +analysis_solved_np = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsSourceSolved, + use_jax=False, +) +fit_solved_np = analysis_solved_np.fit_from(instance=instance_solved) +log_likelihood_solved_np = float(fit_solved_np.log_likelihood) +print("NumPy fit.log_likelihood (solved):", log_likelihood_solved_np) + +EXPECTED_EAGER_LOG_LIKELIHOOD_SOURCE_PLANE_SOLVED = -94.70750992850252 + +np.testing.assert_allclose( + log_likelihood_solved_np, + EXPECTED_EAGER_LOG_LIKELIHOOD_SOURCE_PLANE_SOLVED, + rtol=1e-4, + err_msg=( + f"point_source/source_plane: regression — eager log_likelihood (solved) " + f"drifted (got {log_likelihood_solved_np}, expected " + f"{EXPECTED_EAGER_LOG_LIKELIHOOD_SOURCE_PLANE_SOLVED})" + ), +) +print( + f"Eager regression assertion PASSED (solved): log_likelihood matches " + f"{EXPECTED_EAGER_LOG_LIKELIHOOD_SOURCE_PLANE_SOLVED}" +) + +# NumPy-vs-JAX vmap parity: the analytic solved fit removes the free-cosmology +# vs no-cosmology split as a large source of eager/vmap mismatch — check the +# two paths agree directly, not just against independent literals. +np.testing.assert_allclose( + log_likelihood_solved_np, + float(result_solved[0]), + rtol=1e-4, + err_msg="point_source/source_plane: solved eager vs vmap parity mismatch", +) + +analysis_solved_jit = al.AnalysisPoint( + dataset=dataset, + solver=solver, + fit_positions_cls=al.FitPositionsSourceSolved, + use_jax=True, +) +fit_solved_jit_fn = jax.jit(analysis_solved_jit.fit_from) + +full_pipeline_jits_solved = False +try: + fit_solved = fit_solved_jit_fn(instance_solved) + print("JIT fit.log_likelihood (solved):", fit_solved.log_likelihood) + assert isinstance( + fit_solved.log_likelihood, jnp.ndarray + ), f"expected jax.Array, got {type(fit_solved.log_likelihood)}" + np.testing.assert_allclose( + float(fit_solved.log_likelihood), log_likelihood_solved_np, rtol=1e-4 + ) + full_pipeline_jits_solved = True + print("PASS: jit(fit_from) round-trip matches NumPy scalar (solved).") +except TypeError as e: + print( + "\nBLOCKER: source-plane jit(fit_from) (solved) is gated by:\n" + f" {type(e).__name__}: {e}\n" + " fit_from returns a bare PointSolver at output component [1][1],\n" + " which is not pytree-registered under jax.jit. Tracked in\n" + " PyAutoPrompt/autolens/fit_point_pytree.md.\n" " Eager NumPy regression assertion still PASSED above." ) diff --git a/scripts/point_source/simulators/simple.py b/scripts/point_source/simulators/simple.py index 7ba39201..ec865b7d 100644 --- a/scripts/point_source/simulators/simple.py +++ b/scripts/point_source/simulators/simple.py @@ -6,7 +6,10 @@ function tests in this folder. The dataset is saved to `dataset/point_source/simple/` as -`point_dataset_positions_only.json`. +`point_dataset_positions_only.json` (positions only, used by most scripts in +this folder) and `point_dataset_with_fluxes_and_time_delays.json` (adds +deterministic fluxes and time delays, used by +`jax_likelihood/fluxes_time_delays.py`). __Model__ @@ -74,4 +77,51 @@ file_path=path.join(dataset_path, "tracer.json"), ) +""" +__Fluxes__ + +The flux of each multiple image is the source's total flux multiplied by the magnification at that +image-plane position. Deterministic: no noise is added to the values themselves, only a fixed +fractional noise-map is attached (this simulator must reproduce value-identical output on every +run; note the JSON serializer does not guarantee stable key order between runs). +""" +flux = 1.0 + +magnifications = al.LensCalc.from_tracer(tracer=tracer).magnification_2d_via_hessian_from( + grid=positions +) + +fluxes = [flux * np.abs(magnification) for magnification in magnifications] +fluxes = al.ArrayIrregular(values=fluxes) + +fluxes_noise_map = al.ArrayIrregular(values=[0.1 * f for f in fluxes]) + +""" +__Time Delays__ + +Time delays follow directly from the mass model and image-plane positions (`Tracer.time_delays_from`). +Deterministic for the same reason: a fixed absolute noise-map, no random scatter on the values. +""" +time_delays = tracer.time_delays_from(grid=positions) + +time_delays_noise_map = al.ArrayIrregular(values=[0.5 for _ in time_delays]) + +""" +__Point Dataset (Positions + Fluxes + Time Delays)__ +""" +dataset_with_fluxes_and_time_delays = al.PointDataset( + name="point_0", + positions=positions, + positions_noise_map=0.005, + fluxes=fluxes, + fluxes_noise_map=fluxes_noise_map, + time_delays=time_delays, + time_delays_noise_map=time_delays_noise_map, +) + +al.output_to_json( + obj=dataset_with_fluxes_and_time_delays, + file_path=path.join(dataset_path, "point_dataset_with_fluxes_and_time_delays.json"), +) + print(f"Saved point-source dataset with {len(positions)} images to {dataset_path}")