diff --git a/.gitignore b/.gitignore index bdceaed..e65f73a 100644 --- a/.gitignore +++ b/.gitignore @@ -28,3 +28,6 @@ jax_profiling/dataset/point_source/simple/point_dataset_positions_only.csv # Auto-simulated datacube outputs from datacube simulator — regenerated each run dataset/datacube/sim_*/ + +# Auto-simulated by plotting_alignment/simulator.py — regenerate, don't commit. +dataset/imaging/plotting_alignment/ diff --git a/jax_profiling/gradient/README.md b/jax_profiling/gradient/README.md index 42791a3..f3fa6d9 100644 --- a/jax_profiling/gradient/README.md +++ b/jax_profiling/gradient/README.md @@ -25,6 +25,10 @@ finite and non-zero. | Imaging, `RectangularAdaptDensity` (pixelization over-sampling 1) | **autodiff correct — likelihood is a staircase in mass/shear** | **yes** (see below) | lens-light params FD-matched (≤2e-8); mass/shear: LL is *bit-identical* under ≤1e-6 parameter shifts, so AD's ~zero is the true a.e. derivative and larger-step FD only measures discrete rank-reordering jumps. Gradient-based **mass** inference impossible in this config — a likelihood-design property, not an AD bug (see "Rectangular adaptive mesh" section) | | Imaging, `RectangularAdaptDensity` (pixelization over-sampling 4) | **works** | **yes** (2026-07-09: all 14 params live, AD ≈ FD(h=1e-7) ≤ ~3%, FD converges to AD as h→0; variant D of `jax_grad/imaging_pixelization.py`) | sub-pixel strain between interp queries and knots carries smooth mass information; residual AD-FD gap is micro-staircase contamination of FD, not AD error | | Imaging, `RectangularAdaptImage` + `reg.Adapt` + adapt images + border relocator (pixelization over-sampling 4 — **the production config**) | **works** | **yes** (2026-07-09: all 14 params live, AD ≈ FD(h=1e-7) ≤ ~1%, lens light to 6 digits; variant C of `jax_grad/imaging_pixelization.py`) | full production shape incl. `AdaptImages` weight map and border relocator; mixed precision off in the test (FD needs float64) | +| Imaging, `RectangularKernelAdaptDensity` (os_pix=1, bandwidth=0.1) | **works — the os_pix=1 staircase corner is fixed** | **yes** (2026-07-10: all 14 params incl. mass/shear at STRICT tolerances via FD-step-sweep; variant E of `jax_grad/imaging_pixelization.py`) | kernel-density CDF mesh (PyAutoArray#374): no ranks/sorts, C^∞ transform; FoM *beats* the linear `AdaptDensity` by 4.1e-2 relative at bandwidth 0.1 (value-parity is undefined at os_pix=1 — the linear mesh there is the staircase) | +| Imaging, `RectangularKernelAdaptDensity` (os_pix=4) | **works** | **yes** (2026-07-10: strict FD-step-sweep, all 14 params; variant F) | FoM parity with linear `AdaptDensity` = 2.7e-5 relative at default bandwidth | +| Imaging, `RectangularKernelAdaptImage` + `reg.Adapt` + adapt images + border relocator (os_pix=4, bandwidth=0.1) | **works** | **yes** (2026-07-10: strict FD-step-sweep, all 14 params; variant G) | full production shape; FoM parity floor 6.3e-4 relative (intrinsic: the kernel smooths the adapt-image weights over its bandwidth; swept bandwidth×n_knots 2026-07-10) | +| Interferometer, `RectangularKernelAdaptDensity` + `reg.Adapt` (sparse operator) | **works — the interferometer staircase now has its escape hatch** | **yes** (2026-07-10: strict FD-step-sweep, all 7 mass/shear params live; variant D of `jax_grad/interferometer.py`) | the production-shape adaptive mesh with usable gradients on the no-over-sampling path; FoM parity 3.9e-5 relative at default bandwidth | | Imaging, Delaunay pixelization | **hard error** (re-confirmed 2026-07-09: 3 PASS / 8 ERROR) | n/a | `jax.pure_callback` → `scipy.spatial.Delaunay` has no JVP rule (`PyAutoArray .../interpolator/delaunay.py:126`); see "Why Delaunay gradients are infeasible today" below (phase 2) | | Point source, source-plane χ² (`FitPositionsSource`) | **works** (probe 4/4 PASS; forward `jax.jit` still blocked by the `Grid2DIrregular` xp gap) | **yes** (2026-07-09, rel err ≤ 5e-6; `jax_grad/point_source.py`) | includes magnification-via-Hessian term (3rd derivatives of the potential); flux/H0 legitimately zero in positions-only fits | | Point source, image-plane (`FitPositionsImagePairAll`) | prior probe: **not differentiable** | n/a | `PointSolver` triangle-tiling forward solve uses `jnp.where` masking + integer neighbour lookups | @@ -78,20 +82,35 @@ interpolator refactor). On the refactored main there is no explosion and no micro-staircase contamination of the finite differences, not an autodiff error. - The uniform mesh has no transform and is exactly smooth-differentiable. -**Implications**: (1) HMC/NUTS over mass parameters works with `RectangularUniform` -or with either adaptive mesh at pixelization over-sampling > 1 — which is the -imaging production configuration (`RectangularAdaptImage`, os_pix=4, validated -≤ ~1%); only the os_pix=1 adaptive corner is unusable. **Interferometer is that -corner by construction**: its pixelization has no over-sampling, so the adaptive -meshes carry no smooth mass information there at all (confirmed on the sparse- -operator production path 2026-07-09) — interferometer gradient work must use -`RectangularUniform` (validated) until a smooth transform exists. A continuous -density transform (the paper's actual formulation — CDF from a *smooth* -GP-weighted density rather than empirical point ranks) would remove the -micro-staircase entirely and restore the adaptive meshes everywhere; -(2) even non-gradient samplers see a micro-staircase surface with os_pix=1 — -worth keeping in mind for evidence estimates; (3) PR #281's fix is moot on the -refactored code — do not re-land it. +**Implications** *(updated 2026-07-10 — the continuous transform now exists)*: +(1) HMC/NUTS over mass parameters works with `RectangularUniform`, with either +linear adaptive mesh at pixelization over-sampling > 1 (imaging production +config validated ≤ ~1%), **and now with the kernel-CDF meshes +(`RectangularKernelAdaptDensity` / `RectangularKernelAdaptImage`, +PyAutoArray#374) in every configuration** — the continuous density transform +anticipated below shipped as opt-in mesh classes and is FD-certified at strict +tolerances at os_pix=1, os_pix=4 and on the interferometer sparse path +(variants E/F/G of `imaging_pixelization.py`, variant D of `interferometer.py`). +The linear meshes' staircase documentation above remains accurate — they are +unchanged. (2) even non-gradient samplers see a micro-staircase surface with +the *linear* meshes at os_pix=1 — worth keeping in mind for evidence estimates; +the kernel meshes remove this too. (3) PR #281's fix is moot on the refactored +code — do not re-land it. + +**Kernel-mesh caveats** (2026-07-10): (a) the exact kernel forward is O(M×N) in +memory (over-sampled queries × traced points) — fine at the jax_test/certification +scales, but a 15k-pixel imaging config at os_pix=4 allocates ~60 GB; production +imaging use needs a chunked evaluation (follow-up). The interferometer sparse +path (M = N) is unaffected. (b) FD probing of any pixelized-source likelihood is +poisoned pseudo-randomly by measure-thin solver branch flips (width < 1e-15 in +the parameter, ΔLL ~1.6e-3 on the 8×8 interferometer config up to ~14 on 28×28 +imaging; present under `reg.Constant` and `reg.Adapt`; pre-existing, exposed — +not caused — by the kernel meshes making mass/shear FD certifiable at all). +`jax_grad/util.py` therefore runs an FD-step-sweep for the kernel variants: +per parameter, FD at rel steps {1e-8, 1e-7, 1e-6}, compared at the step closest +to autodiff — clean steps converge to AD at 1e-6..1e-9 relative, so a wrong AD +fails every step. The flips deserve their own investigation (likely +positive-only-solver / PDIP tie-breaks — NNLS-ledger territory). ## Why Delaunay gradients are infeasible today (phase 2a, 2026-07-09) @@ -135,6 +154,14 @@ frozen-triangulation gradients the same way as the rectangular mesh. ## Findings log +- **2026-07-10** (kernel-CDF certification, PyAutoArray#373/#374): the + kernel-density CDF meshes pass strict FD on ALL parameters in every + configuration, including the two previously-dead corners (imaging os_pix=1, + interferometer sparse). Two new cross-cutting findings: (i) measure-thin + solver branch flips poison single-step FD pseudo-randomly (see "Kernel-mesh + caveats" above; `util.compare_gradients` grew an FD-step-sweep mode); + (ii) the exact kernel forward is O(M×N) memory — production-scale imaging + needs chunking before the kernel meshes leave opt-in status. - **2026-07-09** (phase 1, setup): the `jax_grad` FD test scripts previously pointed at datasets that no longer exist (`imaging/simple`, `imaging/source_complex`) while their auto-simulate fallback writes diff --git a/plotting_alignment/kernel_cdf_alignment.py b/plotting_alignment/kernel_cdf_alignment.py new file mode 100644 index 0000000..ce5ec4d --- /dev/null +++ b/plotting_alignment/kernel_cdf_alignment.py @@ -0,0 +1,309 @@ +""" +Kernel-CDF Mesh Alignment Verification +====================================== + +Verifies the rect-adapt plot fix (PyAutoArray#372 / PR#375) end-to-end and +extends it to the kernel-CDF meshes (PyAutoArray#373 / PR#374): renders the +source-plane reconstruction of the ``mass_centre_source_up_more`` dataset +(truth: source Sersic at (y, x) = (1.0, 0.0)) through the real plotting +pipeline (``aaplt.plot_inversion_reconstruction`` — imshow for uniform, +pcolormesh via ``edges_transformed`` for adaptive rectangular, tripcolor for +Delaunay) for every mesh flavour, and reports the bright-core flux centroid +and brightest-cell position of exactly the geometry each renderer draws. + +Pass criterion (issue #372's complaint): the adaptive-mesh pcolormesh +centroids agree with the trusted references (uniform imshow, Delaunay +tripcolor) and the true source position to well under one mesh pixel — the +pre-fix uniform [0, 1] edge partition was off by up to ~1.5 mesh pixels in y. + +Run ``plotting_alignment/simulator.py`` first if the dataset is absent. +""" +from pathlib import Path +import numpy as np + +import autofit as af +import autolens as al +import autolens.plot as aplt +import autoarray.plot as aaplt + +TRUTH = (1.0, 0.0) + +dataset_name = "mass_centre_source_up_more" +dataset_path = Path("dataset", "imaging", "plotting_alignment", dataset_name) +output_path = Path("plotting_alignment", "output", "kernel_cdf") +output_path.mkdir(parents=True, exist_ok=True) + +sub_size = 4 +pixel_scale = 0.05 +mesh_shape = (30, 30) + +dataset = al.Imaging.from_fits( + data_path=dataset_path / "data.fits", + psf_path=dataset_path / "psf.fits", + noise_map_path=dataset_path / "noise_map.fits", + pixel_scales=pixel_scale, + over_sample_size_lp=sub_size, + over_sample_size_pixelization=1, +) +mask = al.Mask2D.circular( + shape_native=dataset.shape_native, pixel_scales=dataset.pixel_scales, radius=3.0 +) +dataset = dataset.apply_mask(mask=mask) +over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from( + grid=dataset.grid, + sub_size_list=[4, 2, 1], + radial_list=[0.3, 0.6], + centre_list=[(0.0, 0.0)], +) +dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size) + +""" +__Model__ + +Lens fixed at the simulator truth (Isothermal ER=1.6, q=0.8 in ell_comps_0, +zero shear) so the reconstruction sits at the true source position. +""" + + +def model_for(mesh): + mass = af.Model(al.mp.Isothermal) + mass.centre.centre_0 = 0.0 + mass.centre.centre_1 = 0.0 + mass.einstein_radius = 1.6 + mass.ell_comps.ell_comps_0 = 0.1111111111111111 + mass.ell_comps.ell_comps_1 = 0.0 + shear = af.Model(al.mp.ExternalShear) + shear.gamma_1 = 0.0 + shear.gamma_2 = 0.0 + lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear) + pixelization = al.Pixelization( + mesh=mesh, regularization=al.reg.Constant(coefficient=1.0) + ) + source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization) + return af.Collection(galaxies=af.Collection(lens=lens, source=source)) + + +def fit_for(mesh): + model = model_for(mesh) + analysis = al.AnalysisImaging( + dataset=dataset, + adapt_images=al.AdaptImages( + galaxy_name_image_dict={ + "('galaxies', 'lens')": dataset.data, + "('galaxies', 'source')": dataset.data, + } + ), + raise_inversion_positions_likelihood_exception=False, + ) + instance = model.instance_from_prior_medians() + return analysis.fit_from(instance) + + +def bright_stats(centres_yx, values): + """Flux centroid over the bright core (>= 50% of max) + argmax position, + computed from exactly the (position, value) pairs the renderer draws.""" + v = np.maximum(np.asarray(values), 0.0) + core = v >= 0.5 * v.max() + centroid = (v[core, None] * centres_yx[core]).sum(axis=0) / v[core].sum() + return centroid, centres_yx[np.argmax(v)] + + +results = {} + +for label, mesh in [ + ("uniform", al.mesh.RectangularUniform(shape=mesh_shape)), + ( + "adapt_image", + al.mesh.RectangularAdaptImage(shape=mesh_shape, weight_power=1.0), + ), + ("adapt_density", al.mesh.RectangularAdaptDensity(shape=mesh_shape)), + ( + "kernel_adapt_density", + al.mesh.RectangularKernelAdaptDensity(shape=mesh_shape, bandwidth=0.1), + ), + ( + "kernel_adapt_image", + al.mesh.RectangularKernelAdaptImage( + shape=mesh_shape, weight_power=1.0, bandwidth=0.1 + ), + ), +]: + fit = fit_for(mesh) + mapper = fit.inversion.cls_list_from(cls=al.Mapper)[0] + reconstruction = np.asarray(fit.inversion.reconstruction_dict[mapper]) + + aaplt.plot_inversion_reconstruction( + pixel_values=reconstruction, + mapper=mapper, + title=f"{label} (truth at {TRUTH})", + zoom_to_brightest=False, + grid=np.array([TRUTH]), + output_path=str(output_path), + output_filename=label, + output_format="png", + ) + + # Value positions exactly as rendered: uniform imshow draws on the mesh + # grid; adaptive pcolormesh draws cells bounded by edges_transformed. + if label == "uniform": + centres = np.asarray(mapper.source_plane_mesh_grid) + else: + y_edges, x_edges = np.asarray(mapper.mesh_geometry.edges_transformed).T + y_mid = 0.5 * (y_edges[:-1] + y_edges[1:]) + x_mid = 0.5 * (x_edges[:-1] + x_edges[1:]) + yy, xx = np.meshgrid(y_mid, x_mid, indexing="ij") + centres = np.stack([yy.ravel(), xx.ravel()], axis=1) + + centroid, peak = bright_stats(centres, reconstruction) + results[label] = (centroid, peak, float(fit.figure_of_merit)) + print( + f"{label:>22}: centroid=({centroid[0]:+.4f}, {centroid[1]:+.4f}) " + f"peak_cell=({peak[0]:+.4f}, {peak[1]:+.4f}) FoM={fit.figure_of_merit:.2f}" + ) + +# Delaunay tripcolor reference (values drawn at the mesh points themselves). +try: + fit = fit_for(al.mesh.Delaunay(pixels=900)) + mapper = fit.inversion.cls_list_from(cls=al.Mapper)[0] + reconstruction = np.asarray(fit.inversion.reconstruction_dict[mapper]) + aaplt.plot_inversion_reconstruction( + pixel_values=reconstruction, + mapper=mapper, + title=f"delaunay (truth at {TRUTH})", + zoom_to_brightest=False, + grid=np.array([TRUTH]), + output_path=str(output_path), + output_filename="delaunay", + output_format="png", + ) + centres = np.asarray(mapper.source_plane_mesh_grid) + centroid, peak = bright_stats(centres, reconstruction) + results["delaunay"] = (centroid, peak, float(fit.figure_of_merit)) + print( + f"{'delaunay':>22}: centroid=({centroid[0]:+.4f}, {centroid[1]:+.4f}) " + f"peak_cell=({peak[0]:+.4f}, {peak[1]:+.4f}) FoM={fit.figure_of_merit:.2f}" + ) +except Exception as e: + print(f"delaunay reference skipped: {type(e).__name__}: {e}") + +""" +__Verdict — node containment + mapper-faithfulness__ + +Two checks per adaptive mesh, splitting "plot geometry" from "mapper physics": + +A. **Containment** — every interior node (where the mapper places value (r, c)) + must lie inside its drawn pcolormesh cell. Post-#375 the node is the exact + U-space midpoint of its cell edges, so a monotone CDF guarantees this; the + check catches any convention regression for all three CDF flavours. + +B. **Mapper-faithfulness** — scatter one over-sampled point's bilinear weights + and compare the drawn centroid (edge-midpoint cells) against what the + mapper actually encodes, sum(w * node position). These must agree to a + fraction of the local cell; the residual |sum(w * node) - query| is the + mapper's own bilinear-in-U interpolation displacement on non-uniform cells + (physics of adaptive meshes, identical pre/post fix) and is reported, not + asserted. +""" +from autoarray.inversion.mesh.interpolator.rectangular import ( + adaptive_rectangular_transformed_grid_from, +) +from autoarray.inversion.mesh.interpolator.rectangular_kernel import ( + adaptive_rectangular_transformed_grid_from_kernel, +) + +n_y, n_x = mesh_shape + +for label, mesh in [ + ("adapt_image", al.mesh.RectangularAdaptImage(shape=mesh_shape, weight_power=1.0)), + ("adapt_density", al.mesh.RectangularAdaptDensity(shape=mesh_shape)), + ( + "kernel_adapt_density", + al.mesh.RectangularKernelAdaptDensity(shape=mesh_shape, bandwidth=0.1), + ), + ( + "kernel_adapt_image", + al.mesh.RectangularKernelAdaptImage( + shape=mesh_shape, weight_power=1.0, bandwidth=0.1 + ), + ), +]: + fit = fit_for(mesh) + mapper = fit.inversion.cls_list_from(cls=al.Mapper)[0] + interpolator = mapper.interpolator + + # Node positions per the mapper convention, through the same CDF machinery. + U_y_nodes = (n_y - np.arange(n_y) - 1.0) / (n_y - 3) + U_x_nodes = (np.arange(n_x) - 1.0) / (n_x - 3) + U_nodes = np.stack([U_y_nodes, U_x_nodes]).T + data_grid = np.asarray(interpolator.data_grid.array) + if "kernel" in label: + nodes_t = adaptive_rectangular_transformed_grid_from_kernel( + data_grid, + U_nodes, + mesh_pixels=n_y, + mesh_weight_map=interpolator.mesh_weight_map, + bandwidth=interpolator.bandwidth, + n_knots=interpolator.n_knots, + ) + else: + nodes_t = adaptive_rectangular_transformed_grid_from( + data_grid, U_nodes, mesh_weight_map=interpolator.mesh_weight_map + ) + node_y, node_x = np.asarray(nodes_t).T + + y_edges, x_edges = np.asarray(mapper.mesh_geometry.edges_transformed).T + + # A: interior nodes inside their drawn cells. Guard-ring cells clamp to + # the data span and collapse to zero width with node == edge — allow a + # few-ulp tolerance for exactly those float ties. + tol = 1e-9 + ok_y = [ + min(y_edges[r], y_edges[r + 1]) - tol + <= node_y[r] + <= max(y_edges[r], y_edges[r + 1]) + tol + for r in range(1, n_y - 1) + ] + ok_x = [ + min(x_edges[c], x_edges[c + 1]) - tol + <= node_x[c] + <= max(x_edges[c], x_edges[c + 1]) + tol + for c in range(1, n_x - 1) + ] + assert all(ok_y) and all(ok_x), f"{label}: a node fell outside its drawn cell" + + # B: drawn centroid vs mapper-encoded centroid for a scattered delta. + positions = np.asarray(interpolator.data_grid.over_sampled.array) + k = int(np.argmin(np.hypot(positions[:, 0] - TRUTH[0], positions[:, 1] - TRUTH[1]))) + query = positions[k] + mappings, _, weights = interpolator._mappings_sizes_weights + idx = np.asarray(mappings)[k] + w = np.asarray(weights)[k] + + y_mid = 0.5 * (y_edges[:-1] + y_edges[1:]) + x_mid = 0.5 * (x_edges[:-1] + x_edges[1:]) + rows, cols = idx // n_x, idx % n_x + drawn = np.array( + [(w * y_mid[rows]).sum() / w.sum(), (w * x_mid[cols]).sum() / w.sum()] + ) + encoded = np.array( + [(w * node_y[rows]).sum() / w.sum(), (w * node_x[cols]).sum() / w.sum()] + ) + cell = max( + np.abs(np.diff(y_edges))[rows].max(), np.abs(np.diff(x_edges))[cols].max() + ) + d_faithful = np.hypot(*(drawn - encoded)) + d_interp = np.hypot(*(encoded - query)) + print( + f"{label:>22}: drawn-vs-encoded = {d_faithful:.4f}\" " + f"({d_faithful / cell:.2f} cells) | mapper interp displacement = " + f"{d_interp:.4f}\" ({d_interp / cell:.2f} cells; physics, unasserted)" + ) + assert d_faithful < 0.5 * cell, ( + f"{label}: drawn centroid deviates from the mapper-encoded centroid by " + f"{d_faithful:.4f}\" (> 0.5 local cell) — plot geometry unfaithful." + ) + +print( + "\nkernel_cdf_alignment.py: all nodes contained in their drawn cells and " + "the plot is faithful to the mapper on every adaptive mesh (post-#375 edges)." +) diff --git a/plotting_alignment/simulator.py b/plotting_alignment/simulator.py index 9817105..3e8b7bd 100644 --- a/plotting_alignment/simulator.py +++ b/plotting_alignment/simulator.py @@ -162,10 +162,19 @@ Output the simulated dataset to the dataset path as .fits files. """ -dataset.output_to_fits( - data_path=dataset_path / "data.fits", - psf_path=dataset_path / "psf.fits", - noise_map_path=dataset_path / "noise_map.fits", +al.output_to_fits( + values=dataset.data.native, + file_path=dataset_path / "data.fits", + overwrite=True, +) +al.output_to_fits( + values=dataset.psf.kernel.native, + file_path=dataset_path / "psf.fits", + overwrite=True, +) +al.output_to_fits( + values=dataset.noise_map.native, + file_path=dataset_path / "noise_map.fits", overwrite=True, ) @@ -179,9 +188,17 @@ Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files. """ -aplt.plot_array(array=dataset.data, output=aplt.Output(path=dataset_path, format="png")) +aplt.plot_array( + array=dataset.data, + output_path=str(dataset_path), + output_filename="data", + output_format="png", +) aplt.subplot_tracer( - tracer=tracer, grid=grid, output=aplt.Output(path=dataset_path, format="png") + tracer=tracer, + grid=grid, + output_path=str(dataset_path), + output_format="png", ) """