diff --git a/scripts/jax_grad/interferometer.py b/scripts/jax_grad/interferometer.py new file mode 100644 index 00000000..bbc8af49 --- /dev/null +++ b/scripts/jax_grad/interferometer.py @@ -0,0 +1,283 @@ +""" +Tests JAX gradients of the interferometer log-likelihood, in two stages +(finiteness, then autodiff-vs-central-finite-difference correctness — see +``util.py``), for the two configurations used in practice: + +**Variant A — parametric light profiles** (``lp.Sersic`` standard + +``lp_linear.Sersic``): the visibility-space analogue of +``imaging_lp.py``. The source is the only light in interferometer data, so +the model has no lens light component and the evaluation point is anchored +near the simulator truth so the positive-only NNLS keeps the linear source +live. + +**Variant B — ``RectangularAdaptDensity`` via the sparse linear-algebra +path**: mirrors ``jax_likelihood_functions/interferometer/rectangular_sparse.py`` +exactly — ``TransformerDFT`` + ``apply_sparse_operator(use_jax=True)`` (the +sparse NUFFT response-matrix formalism; the operator is aux state built once +outside the JIT trace), ``RectangularAdaptDensity`` mesh + ``reg.Adapt()`` + +``al.AdaptImages``. **Measured verdict (2026-07-09): the imaging os_pix=1 +staircase applies** — interferometer pixelization has no over-sampling, so the +mesh's rank-transform queries coincide with its knots and the likelihood is +invariant to smooth mass perturbations: every mass/shear autodiff gradient is +exactly zero (correct — FD shows only ~1e-7-scale micro-jumps from rank +re-orderings, no smooth slope). With the model having no lens light, that means +**no usable gradients at all** in this configuration. The assertions document +this staircase so a change in mesh differentiability fails loudly. + +**Variant C — ``RectangularUniform`` via the same sparse path**: the working +alternative for gradient-based inference — no adaptive transform, so mass/shear +gradients are live and strictly FD-matched. + +See the audit README +(`autolens_workspace_developer/jax_profiling/gradient/README.md`). +""" + +import numpy as np +import jax +import jax.numpy as jnp +from os import path + +import autofit as af +import autolens as al + +import util + +""" +__Mask + Dataset__ +""" +real_space_mask = al.Mask2D.circular( + shape_native=(256, 256), + pixel_scales=0.1, + radius=3.0, +) + +dataset_name = "simple" +dataset_path = path.join("dataset", "interferometer", dataset_name) + +""" +__Dataset Auto-Simulation__ +""" +if al.util.dataset.should_simulate(dataset_path): + import subprocess + import sys + + subprocess.run( + [sys.executable, "scripts/jax_likelihood_functions/interferometer/simulator.py"], + check=True, + ) + +dataset = al.Interferometer.from_fits( + data_path=path.join(dataset_path, "data.fits"), + noise_map_path=path.join(dataset_path, "noise_map.fits"), + uv_wavelengths_path=path.join(dataset_path, "uv_wavelengths.fits"), + real_space_mask=real_space_mask, + transformer_class=al.TransformerDFT, +) + +from autofit.non_linear.fitness import Fitness + + +def mass_and_shear(): + """ + Truth-centred lens mass model (simulator: Isothermal ER=1.6, shear 0.05). + """ + mass = af.Model(al.mp.Isothermal) + mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.005) + mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.005) + mass.einstein_radius = af.GaussianPrior(mean=1.6, sigma=0.05) + mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.05, sigma=0.01) + mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.05, sigma=0.01) + shear = af.Model(al.mp.ExternalShear) + shear.gamma_1 = af.GaussianPrior(mean=0.05, sigma=0.005) + shear.gamma_2 = af.GaussianPrior(mean=0.05, sigma=0.005) + return mass, shear + + +def param_vector_from(model): + param_vector = jnp.array(model.physical_values_from_prior_medians) + key = jax.random.PRNGKey(42) + perturbation = jax.random.uniform( + key, shape=param_vector.shape, minval=0.001, maxval=0.005 + ) + return param_vector + perturbation + + +def finiteness_checks(fitness, param_vector, n_params): + value, grad = jax.value_and_grad(fitness.call)(param_vector) + print(f"Log likelihood = {float(value):.6f}") + assert np.isfinite(float(value)), "Log likelihood is not finite" + assert grad.shape == (n_params,), f"Gradient shape mismatch: {grad.shape}" + assert np.all( + np.isfinite(np.array(grad)) + ), f"Gradient contains non-finite values: {np.array(grad)}" + assert not np.all(np.array(grad) == 0.0), "Gradient is all zeros" + return grad + + +""" +__Variant A: parametric light profiles (standard + linear Sersic source)__ +""" +for variant, bulge_cls in [ + ("lp.Sersic (standard)", al.lp.Sersic), + ("lp_linear.Sersic (linear)", al.lp_linear.Sersic), +]: + print(f"\n=== interferometer {variant} ===") + + source_bulge = af.Model(bulge_cls) + source_bulge.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.005) + source_bulge.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.005) + source_bulge.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.05, sigma=0.01) + source_bulge.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.05, sigma=0.01) + source_bulge.effective_radius = af.GaussianPrior(mean=1.0, sigma=0.05) + source_bulge.sersic_index = af.GaussianPrior(mean=1.0, sigma=0.2) + + mass, shear = mass_and_shear() + lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear) + source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge) + model = af.Collection(galaxies=af.Collection(lens=lens, source=source)) + + analysis = al.AnalysisInterferometer(dataset=dataset) + + fitness = Fitness( + model=model, + analysis=analysis, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99, + ) + + param_vector = param_vector_from(model) + param_names = util.parameter_names_from(model) + + finiteness_checks(fitness, param_vector, n_params=len(param_names)) + + f_jit = jax.jit(fitness.call) + + util.assert_eager_jit_consistent(fitness.call, f_jit, param_vector) + + comparison = util.compare_gradients( + fitness.call, + param_vector, + param_names=param_names, + f_fd=f_jit, + ) + + util.assert_gradients_match(comparison) + + print(f"interferometer {variant}: autodiff matches finite differences.") + +""" +__Variants B + C: rectangular pixelized source via the sparse-operator path__ + +The sparse NUFFT operator is aux state and must be built once, outside any JIT +trace (mirrors ``rectangular_sparse.py``). +""" +dataset_sparse = dataset.apply_sparse_operator(use_jax=True, show_progress=True) + +mesh_shape = (8, 8) + + +def sparse_fitness(mesh, regularization): + mass, shear = mass_and_shear() + lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear) + + pixelization = al.Pixelization(mesh=mesh, regularization=regularization) + source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization) + model = af.Collection(galaxies=af.Collection(lens=lens, source=source)) + + bulge = al.lp.Sersic() + adapt_image = bulge.image_2d_from(grid=dataset_sparse.grid) + adapt_images = al.AdaptImages( + galaxy_name_image_dict={ + "('galaxies', 'lens')": adapt_image, + "('galaxies', 'source')": adapt_image, + } + ) + + analysis = al.AnalysisInterferometer( + dataset=dataset_sparse, + adapt_images=adapt_images, + raise_inversion_positions_likelihood_exception=False, + ) + + fitness = Fitness( + model=model, + analysis=analysis, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99, + ) + return fitness, param_vector_from(model), util.parameter_names_from(model) + + +""" +__Variant B: RectangularAdaptDensity — the documented staircase__ +""" +print("\n=== interferometer RectangularAdaptDensity + reg.Adapt, sparse operator ===") + +fitness, param_vector, param_names = sparse_fitness( + mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape), + regularization=al.reg.Adapt(), +) + +value, grad = jax.value_and_grad(fitness.call)(param_vector) +print(f"Log likelihood = {float(value):.6f}") +assert np.isfinite(float(value)), "Log likelihood is not finite" +assert np.all( + np.isfinite(np.array(grad)) +), f"Gradient contains non-finite values: {np.array(grad)}" + +f_jit = jax.jit(fitness.call) + +util.assert_eager_jit_consistent(fitness.call, f_jit, param_vector) + +# The staircase: with no over-sampling the adaptive mesh's rank transform makes +# the likelihood invariant to smooth mass/shear perturbations, and *every* model +# parameter here is mass/shear (interferometer data has no lens light). The +# correct autodiff gradient is therefore ~zero across the board. If this +# assertion ever fails the mesh has become differentiable — rerun the full FD +# audit and update this script + the audit README. +assert np.all(np.abs(np.array(grad)) < 1e-6), ( + "Autodiff mass/shear gradients are no longer ~zero on the sparse " + f"RectangularAdaptDensity path: {np.array(grad)}" +) + +print( + "interferometer sparse RectangularAdaptDensity: staircase confirmed — " + "all autodiff gradients ~zero (correct; no smooth mass information)." +) + +""" +__Variant C: RectangularUniform — the gradient-capable alternative__ +""" +print("\n=== interferometer RectangularUniform + reg.Constant, sparse operator ===") + +fitness, param_vector, param_names = sparse_fitness( + mesh=al.mesh.RectangularUniform(shape=mesh_shape), + regularization=al.reg.Constant(coefficient=1.0), +) + +grad = finiteness_checks(fitness, param_vector, n_params=len(param_names)) + +f_jit = jax.jit(fitness.call) + +util.assert_eager_jit_consistent(fitness.call, f_jit, param_vector) + +comparison = util.compare_gradients( + fitness.call, + param_vector, + param_names=param_names, + f_fd=f_jit, +) + +util.assert_gradients_match(comparison) + +assert np.all(np.abs(comparison["ad"]) > 0.0), ( + "A parameter gradient is exactly zero on the sparse RectangularUniform " + f"path: {[(n, a) for n, a in zip(param_names, comparison['ad']) if a == 0.0]}" +) + +print( + "interferometer sparse RectangularUniform: all gradients live and " + "FD-matched." +) + +print("\ninterferometer.py JAX gradient checks passed.")