diff --git a/likelihood_breakdown/cluster/image_plane.py b/likelihood_breakdown/cluster/image_plane.py new file mode 100644 index 0000000..249d0dd --- /dev/null +++ b/likelihood_breakdown/cluster/image_plane.py @@ -0,0 +1,384 @@ +""" +JAX Profiling: Cluster Image-Plane Likelihood — Per-Step Breakdown +=================================================================== + +Decomposes the cluster point-source **image-plane chi-squared** likelihood +(``FitPositionsImagePairRepeat``, the model-fit default) into its pipeline +steps for the standard cluster model (2 main dPIE + 10 scaling dPIE + NFW +host at z = 0.5; 2 point sources at z = 1.0 / 2.0 — multi-plane). + +Where the source-plane likelihood (``source_plane.py``) only ray-traces the +observed positions *backwards*, the image-plane likelihood **forward-solves +the lens equation** for every source: tile the image plane in triangles, +trace them to the source plane, keep the ones containing the source centre, +subdivide, repeat to sub-pixel precision. That solve dominates everything +else by orders of magnitude and is the reason cluster image-plane fits need +JAX — this script makes its cost (and its one-off JIT compile cost, which a +sampler pays once but a single fit pays in full) visible per source plane. + +Steps profiled: + +1. Back-trace observed positions → model source centres (setup, eager). +2. Triangle-tiling PointSolver solve, JIT-compiled per source plane — the + dominant step. Lower/compile/first-call/steady-state are reported + separately so compile amortisation is explicit. +3. ``FitPositionsImagePairRepeat`` log-likelihood per system (eager, + nearest-pair chi-squared given the solved positions). + +The solver grid below (200x200 @ 0.7", precision 0.01") is the +tutorial-scale configuration of the workspace cluster scripts, chosen so the +one-off compile stays at minutes; production precision (0.001") multiplies +the triangle fan-out, not the structure of the breakdown. + +Output +------ + +Results JSON and PNG are written to ``results/breakdown/cluster/`` using the +basename ``image_plane_breakdown_v{al_version}``. +""" + +# --------------------------------------------------------------------------- +# AUTOLENS_PROFILING_SMOKE=1 short-circuit (CI lint smoke). +# --------------------------------------------------------------------------- +import os as _smoke_os +import sys +import time +from contextlib import contextmanager +from pathlib import Path + +import autolens as al +import jax +import jax.numpy as jnp +import numpy as np + +if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": + print(f"[smoke] {__file__}: imports + module setup OK; exiting.") + sys.exit(0) + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) +from _profile_cli import ( # noqa: E402 + auto_simulate_if_missing, + device_info_dict, + parse_profile_cli, + resolve_output_paths, +) + +_cli = parse_profile_cli() + +_script_dir = Path(__file__).resolve().parent +_workspace_root = _script_dir.parents[1] + + +# --------------------------------------------------------------------------- +# Profiling helpers (house pattern) +# --------------------------------------------------------------------------- + + +class Timer: + def __init__(self): + self.records: list[tuple[str, float]] = [] + + @contextmanager + def section(self, label: str): + start = time.perf_counter() + yield + elapsed = time.perf_counter() - start + self.records.append((label, elapsed)) + print(f" [{label}] {elapsed:.4f} s") + + def summary(self): + print("\n" + "=" * 70) + print("PROFILING SUMMARY") + print("=" * 70) + max_label = max(len(r[0]) for r in self.records) + total = 0.0 + for label, elapsed in self.records: + print(f" {label:<{max_label}} {elapsed:>10.4f} s") + total += elapsed + print("-" * 70) + print(f" {'TOTAL':<{max_label}} {total:>10.4f} s") + print("=" * 70) + + +def block(x): + if hasattr(x, "block_until_ready"): + x.block_until_ready() + return x + + +def jit_profile(func, label, *args, n_repeats=10): + jitted = jax.jit(func) + + with timer.section(f"{label}_lower"): + lowered = jitted.lower(*args) + + with timer.section(f"{label}_compile"): + compiled = lowered.compile() + + with timer.section(f"{label}_first_call"): + result = compiled(*args) + block(result) + + with timer.section(f"{label}_steady_x{n_repeats}"): + for _ in range(n_repeats): + result = compiled(*args) + block(result) + + per_call = timer.records[-1][1] / n_repeats + print(f" -> per-call avg: {per_call:.6f} s") + return compiled, result + + +timer = Timer() +likelihood_steps = [] + + +# =================================================================== +# PART A — Setup +# =================================================================== +dataset_path = _workspace_root / "dataset" / "cluster" / "simple" + +auto_simulate_if_missing( + dataset_path, + dataset_type="cluster", + instrument="simple", + workspace_root=_workspace_root, +) + +dataset_list = al.list_from_csv(file_path=dataset_path / "point_datasets.csv") +print(f"Loaded {len(dataset_list)} point-source systems from {dataset_path}.") + +redshift_lens = 0.5 + +main_lens_params = [ + ((0.0, 0.0), 8.0, 20.0, 3.0), + ((10.0, 8.0), 5.0, 12.0, 1.2), +] +main_lens_galaxies = [ + al.Galaxy( + redshift=redshift_lens, + mass=al.mp.dPIEMassSph(centre=centre, ra=ra, rs=rs, b0=b0), + ) + for centre, ra, rs, b0 in main_lens_params +] + +scaling_table = al.galaxy_table_from_csv(file_path=dataset_path / "scaling_galaxies.csv") +SCALING_B0_REF, SCALING_RS_REF, SCALING_RA, SCALING_EXPONENT = 0.12, 10.0, 0.1, 0.5 +_lum_ref = max(scaling_table.luminosities) +scaling_galaxies = [ + al.Galaxy( + redshift=redshift_lens, + mass=al.mp.dPIEMassSph( + centre=tuple(centre), + ra=SCALING_RA, + rs=SCALING_RS_REF * (lum / _lum_ref) ** SCALING_EXPONENT, + b0=SCALING_B0_REF * (lum / _lum_ref) ** SCALING_EXPONENT, + ), + ) + for centre, lum in zip(scaling_table.centres, scaling_table.luminosities) +] + +host_halo_galaxy = al.Galaxy( + redshift=redshift_lens, + dark=al.mp.NFWMCRLudlowSph( + centre=(0.0, 0.0), + mass_at_200=10**15.3, + redshift_object=redshift_lens, + redshift_source=max(float(d.redshift) for d in dataset_list), + ), +) + +source_galaxies = [ + al.Galaxy(redshift=float(d.redshift), **{d.name: al.ps.Point(centre=(0.0, 0.0))}) + for d in dataset_list +] + +tracer = al.Tracer( + galaxies=main_lens_galaxies + scaling_galaxies + [host_halo_galaxy] + source_galaxies +) +n_mass_profiles = len(main_lens_galaxies) + len(scaling_galaxies) + 1 +print(f"Tracer: {len(tracer.planes)} planes, {n_mass_profiles} mass components.") + +positions_list = [np.atleast_2d(np.asarray(d.positions)) for d in dataset_list] +plane_indices = [ + tracer.plane_index_via_redshift_from(redshift=float(d.redshift)) for d in dataset_list +] + +# --------------------------------------------------------------------------- +# Step 1 — model source centres: back-trace the observed images and take the +# per-system centroid (what a model fit derives from its Point centre; here +# the truth-adjacent centroid keeps the solve realistic). Eager + cheap. +# --------------------------------------------------------------------------- +source_centres = [] +with timer.section("step1_source_centres"): + for positions, plane_index in zip(positions_list, plane_indices): + traced = np.asarray( + tracer.traced_grid_2d_list_from(grid=al.Grid2DIrregular(positions))[plane_index] + ) + source_centres.append(tuple(traced.mean(axis=0))) +likelihood_steps.append(("1 source centres (back-trace)", timer.records[-1][1])) + +# --------------------------------------------------------------------------- +# Step 2 — the triangle-tiling PointSolver forward-solve, per source plane. +# Pytree registration mirrors simulators/cluster.py: the model classes are +# registered via autofit's register_model on an af.Model mirror, and Tracer +# itself via register_instance_pytree (cosmology excluded from flattening). +# --------------------------------------------------------------------------- +import autofit as af # noqa: E402 +from autoarray.abstract_ndarray import register_instance_pytree # noqa: E402 +from autofit.jax import register_model as _register_model_pytrees # noqa: E402 +from autolens.lens.tracer import Tracer # noqa: E402 + +_registration_model = af.Collection( + galaxies=af.Collection( + af.Model( + al.Galaxy, + redshift=redshift_lens, + mass=af.Model(al.mp.dPIEMassSph, centre=(0.0, 0.0), ra=1.0, rs=10.0, b0=1.0), + ), + af.Model( + al.Galaxy, + redshift=redshift_lens, + dark=af.Model( + al.mp.NFWMCRLudlowSph, + centre=(0.0, 0.0), + mass_at_200=10**15.3, + redshift_object=redshift_lens, + redshift_source=max(float(d.redshift) for d in dataset_list), + ), + ), + *[ + af.Model( + al.Galaxy, + redshift=float(d.redshift), + **{d.name: af.Model(al.ps.Point, centre=(0.0, 0.0))}, + ) + for d in dataset_list + ], + ) +) +_register_model_pytrees(_registration_model) +register_instance_pytree(Tracer, no_flatten=("cosmology",)) + +solver = al.PointSolver.for_grid( + grid=al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.7), + pixel_scale_precision=0.01, + use_jax=True, +) + +predicted_per_system = [] +for dataset, centre in zip(dataset_list, source_centres): + + def solve(source_plane_coordinate, _z=float(dataset.redshift)): + return solver.solve( + tracer=tracer, + source_plane_coordinate=source_plane_coordinate, + plane_redshift=_z, + ) + + _, predicted = jit_profile(solve, f"step2_solve_{dataset.name}", jnp.array(centre), n_repeats=3) + predicted_per_system.append(predicted) + likelihood_steps.append( + ( + f"2.{dataset.name} PointSolver solve (z={float(dataset.redshift):.1f})", + timer.records[-1][1] / 3, + ) + ) + +# --------------------------------------------------------------------------- +# Step 3 — pairing + chi-squared: the production FitPositionsImagePairRepeat, +# timed eagerly per system (the pairing is trivial next to the solve; the fit +# re-runs the solve internally, so its time is reported as fit-total and the +# pairing overhead is the difference from step 2). +# --------------------------------------------------------------------------- +fit_log_likelihoods = [] +for dataset in dataset_list: + with timer.section(f"step3_fit_total_{dataset.name}"): + fit = al.FitPositionsImagePairRepeat( + name=dataset.name, + data=dataset.positions, + noise_map=dataset.positions_noise_map, + tracer=tracer, + solver=solver, + ) + fit_log_likelihoods.append(float(fit.log_likelihood)) + likelihood_steps.append( + (f"3.{dataset.name} FitPositionsImagePairRepeat (fit total)", timer.records[-1][1]) + ) + +log_likelihood_total = sum(fit_log_likelihoods) +print(f"\n image-plane log likelihood (sum over systems): {log_likelihood_total:.6e}") + + +# =================================================================== +# PART B — Summary + artifacts +# =================================================================== +import json + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +al_version = al.__version__ + +timer.summary() + +print("\n" + "=" * 70) +print(f"PER-STEP BREAKDOWN SUMMARY — CLUSTER IMAGE-PLANE — v{al_version}") +print("=" * 70) +max_label = max(len(label) for label, _ in likelihood_steps) +step_total = 0.0 +for i, (label, per_call) in enumerate(likelihood_steps, 1): + print(f" {i:>2}. {label:<{max_label}} {per_call:>12.6f} s") + step_total += per_call +print("-" * 70) +print(f" {'TOTAL (step-by-step)':<{max_label}} {step_total:>12.6f} s") +print("=" * 70) + +breakdown_summary = { + "autolens_version": al_version, + "device": device_info_dict(), + "configuration": { + "n_systems": len(dataset_list), + "n_images_total": int(sum(len(p) for p in positions_list)), + "n_planes": len(tracer.planes), + "n_mass_profiles": n_mass_profiles, + "likelihood": "image_plane (FitPositionsImagePairRepeat)", + "solver_grid": "200x200 @ 0.7 arcsec", + "solver_pixel_scale_precision": 0.01, + }, + "steps": {label: per_call for label, per_call in likelihood_steps}, + "total_step_by_step": step_total, + "log_likelihood": log_likelihood_total, +} + +dict_path, chart_path = resolve_output_paths( + _cli, + default_dir=_workspace_root / "results" / "breakdown" / "cluster", + default_basename=f"image_plane_breakdown_v{al_version}", +) +dict_path.write_text(json.dumps(breakdown_summary, indent=2)) +print(f"\n Results dict saved to: {dict_path}") + +labels = [label for label, _ in likelihood_steps] +times = [per_call for _, per_call in likelihood_steps] +fig, ax = plt.subplots(figsize=(10, 5)) +y_pos = range(len(labels)) +bars = ax.barh(y_pos, times, color="#4C72B0", edgecolor="white", height=0.6) +for bar, t in zip(bars, times): + ax.text( + bar.get_width() + max(times) * 0.01, + bar.get_y() + bar.get_height() / 2, + f"{t:.6f} s", + va="center", + fontsize=9, + ) +ax.set_yticks(y_pos) +ax.set_yticklabels(labels, fontsize=10) +ax.invert_yaxis() +ax.set_xlabel("Time per call (s)", fontsize=11) +fig.suptitle(f"Cluster image-plane likelihood breakdown — v{al_version}", fontsize=12) +fig.tight_layout() +fig.savefig(chart_path, dpi=150) +print(f" Bar chart saved to: {chart_path}") diff --git a/likelihood_breakdown/cluster/source_plane.py b/likelihood_breakdown/cluster/source_plane.py new file mode 100644 index 0000000..c0e142f --- /dev/null +++ b/likelihood_breakdown/cluster/source_plane.py @@ -0,0 +1,452 @@ +""" +JAX Profiling: Cluster Source-Plane Likelihood — Per-Step Breakdown +==================================================================== + +Decomposes the cluster point-source **source-plane chi-squared** likelihood +(``FitPositionsSource``) into its pipeline steps and JIT-profiles each one +separately, for the standard cluster model: + + - 2 individually-modelled dPIE main lenses + 10 scaling-tier dPIE members + (reference-anchored relation) + 1 NFW host halo, all at z = 0.5; + - 2 point sources at *different* redshifts (z = 1.0, 2.0) — a genuine + multi-plane system, so ray-tracing pays the per-plane recursion. + +This is Lenstool's default likelihood: observed multiple-image positions are +ray-traced *back* to each source's plane and the scatter about the group +centroid is penalised. No lens-equation solve is involved, which is why the +whole thing is orders of magnitude cheaper than the image-plane likelihood +(``image_plane.py`` in this folder profiles that one). + +Steps profiled: + +1. Multi-plane ray-tracing of the observed positions, per source plane — + the deflection stack over all 13 mass profiles is one hot spot. +2. Magnification at every observed position via the tracer Hessian + (``LensCalc.magnification_2d_via_hessian_from``) — the production + source-plane chi-squared weights each residual by its magnification, so + image-plane noise maps correctly into the source plane. This evaluates + the full deflection stack several more times per position and competes + with step 1 for the budget. +3. Magnification-weighted chi-squared per system: the distance of each + traced position to the *model's* source centre (name pairing hands the + fit the ``Point`` profile centre — the barycenter is only a fallback when + the tracer has no matching profile), chi2_i = dist_i^2 * mag_i^2 / sigma_i^2. +4. Total log-likelihood assembly. The noise normalization uses the + magnification-scaled effective noise, sum ln(2 pi (sigma_i / mag_i)^2) — + consistent with the chi-squared's source-plane noise mapping. + +The steps close over the (fixed) tracer rather than passing it as a JIT +pytree argument: each step compiles exactly once here, so registration +buys nothing and the closure keeps the decomposition free of pytree +plumbing. (In a model-fit the tracer changes per call — see +``simulators/cluster.py`` for the pytree-argument pattern used there.) + +Reference check: the summed per-step log-likelihood is asserted against the +eager ``al.FitPositionsSource(profile=None)`` values at ``rtol=1e-4``, so the +decomposition is provably the production calculation. Note the known JAX +constraint: the *end-to-end* source-plane fit is not JIT-compilable today +(see ``autolens_workspace_test/scripts/CLAUDE.md`` — the source-plane entry +is marked JIT-blocked); the per-step decomposition below sidesteps that by +compiling the numerical core of each step in isolation. + +Output +------ + +Results JSON and PNG are written to ``results/breakdown/cluster/`` using the +basename ``source_plane_breakdown_v{al_version}``. +""" + +# --------------------------------------------------------------------------- +# AUTOLENS_PROFILING_SMOKE=1 short-circuit (CI lint smoke): verify the import +# graph + module-level setup succeeded without running the full pipeline. +# --------------------------------------------------------------------------- +import os as _smoke_os +import sys +import time +from contextlib import contextmanager +from pathlib import Path + +import autolens as al +import jax +import jax.numpy as jnp +import numpy as np + +if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": + print(f"[smoke] {__file__}: imports + module setup OK; exiting.") + sys.exit(0) + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) +from _profile_cli import ( # noqa: E402 + auto_simulate_if_missing, + device_info_dict, + parse_profile_cli, + resolve_output_paths, +) + +_cli = parse_profile_cli() + +_script_dir = Path(__file__).resolve().parent +_workspace_root = _script_dir.parents[1] + + +# --------------------------------------------------------------------------- +# Profiling helpers (house pattern — see likelihood_breakdown/imaging/mge.py) +# --------------------------------------------------------------------------- + + +class Timer: + """Accumulates named timing measurements and prints a summary.""" + + def __init__(self): + self.records: list[tuple[str, float]] = [] + + @contextmanager + def section(self, label: str): + start = time.perf_counter() + yield + elapsed = time.perf_counter() - start + self.records.append((label, elapsed)) + print(f" [{label}] {elapsed:.4f} s") + + def summary(self): + print("\n" + "=" * 70) + print("PROFILING SUMMARY") + print("=" * 70) + max_label = max(len(r[0]) for r in self.records) + total = 0.0 + for label, elapsed in self.records: + print(f" {label:<{max_label}} {elapsed:>10.4f} s") + total += elapsed + print("-" * 70) + print(f" {'TOTAL':<{max_label}} {total:>10.4f} s") + print("=" * 70) + + +def block(x): + if hasattr(x, "block_until_ready"): + x.block_until_ready() + return x + + +def jit_profile(func, label, *args, n_repeats=10): + """JIT-compile *func*, time lower/compile/first-call/steady-state.""" + jitted = jax.jit(func) + + with timer.section(f"{label}_lower"): + lowered = jitted.lower(*args) + + with timer.section(f"{label}_compile"): + compiled = lowered.compile() + + with timer.section(f"{label}_first_call"): + result = compiled(*args) + block(result) + + with timer.section(f"{label}_steady_x{n_repeats}"): + for _ in range(n_repeats): + result = compiled(*args) + block(result) + + per_call = timer.records[-1][1] / n_repeats + print(f" -> per-call avg: {per_call:.6f} s") + return compiled, result + + +timer = Timer() +likelihood_steps = [] # (label, per_call_seconds) + + +# =================================================================== +# PART A — Setup (not JIT-compiled) +# =================================================================== + +# --------------------------------------------------------------------------- +# 1. Dataset (auto-simulated on first run via simulators/cluster.py) +# --------------------------------------------------------------------------- +dataset_path = _workspace_root / "dataset" / "cluster" / "simple" + +auto_simulate_if_missing( + dataset_path, + dataset_type="cluster", + instrument="simple", + workspace_root=_workspace_root, +) + +dataset_list = al.list_from_csv(file_path=dataset_path / "point_datasets.csv") +print(f"Loaded {len(dataset_list)} point-source systems from {dataset_path}.") + +# --------------------------------------------------------------------------- +# 2. Tracer — the standard cluster model, at the simulator truth values. +# Constants mirror simulators/cluster.py (which mirrors the workspace +# cluster simulator + its reference-anchored scaling relation). +# --------------------------------------------------------------------------- +redshift_lens = 0.5 + +main_lens_params = [ + ((0.0, 0.0), 8.0, 20.0, 3.0), + ((10.0, 8.0), 5.0, 12.0, 1.2), +] +main_lens_galaxies = [ + al.Galaxy( + redshift=redshift_lens, + mass=al.mp.dPIEMassSph(centre=centre, ra=ra, rs=rs, b0=b0), + ) + for centre, ra, rs, b0 in main_lens_params +] + +scaling_table = al.galaxy_table_from_csv(file_path=dataset_path / "scaling_galaxies.csv") +SCALING_B0_REF, SCALING_RS_REF, SCALING_RA, SCALING_EXPONENT = 0.12, 10.0, 0.1, 0.5 +_lum_ref = max(scaling_table.luminosities) +scaling_galaxies = [ + al.Galaxy( + redshift=redshift_lens, + mass=al.mp.dPIEMassSph( + centre=tuple(centre), + ra=SCALING_RA, + rs=SCALING_RS_REF * (lum / _lum_ref) ** SCALING_EXPONENT, + b0=SCALING_B0_REF * (lum / _lum_ref) ** SCALING_EXPONENT, + ), + ) + for centre, lum in zip(scaling_table.centres, scaling_table.luminosities) +] + +host_halo_galaxy = al.Galaxy( + redshift=redshift_lens, + dark=al.mp.NFWMCRLudlowSph( + centre=(0.0, 0.0), + mass_at_200=10**15.3, + redshift_object=redshift_lens, + redshift_source=max(float(d.redshift) for d in dataset_list), + ), +) + +# Point centres at the simulator truth (simulators/cluster.py source_centres) — +# name pairing hands these to the fit as the source-plane reference points. +TRUTH_SOURCE_CENTRES = [(0.3, 0.5), (-0.8, 1.2)] +source_galaxies = [ + al.Galaxy(redshift=float(d.redshift), **{d.name: al.ps.Point(centre=centre)}) + for d, centre in zip(dataset_list, TRUTH_SOURCE_CENTRES) +] + +tracer = al.Tracer( + galaxies=main_lens_galaxies + scaling_galaxies + [host_halo_galaxy] + source_galaxies +) +n_mass_profiles = len(main_lens_galaxies) + len(scaling_galaxies) + 1 +print( + f"Tracer: {len(tracer.planes)} planes, {n_mass_profiles} mass components " + f"(2 main dPIE + {len(scaling_galaxies)} scaling dPIE + 1 NFW host)." +) + +positions_list = [np.atleast_2d(np.asarray(d.positions)) for d in dataset_list] +noise_list = [np.asarray(d.positions_noise_map) for d in dataset_list] +plane_indices = [ + tracer.plane_index_via_redshift_from(redshift=float(d.redshift)) for d in dataset_list +] + + +# =================================================================== +# PART B — Per-step JIT profiling +# =================================================================== + +# --------------------------------------------------------------------------- +# Step 1 — multi-plane ray-tracing of the observed positions, per system. +# The recursion walks every plane up to the source's, applying the scaled +# deflections of all 13 mass profiles; this is where the cluster's many- +# profile cost lives. +# --------------------------------------------------------------------------- +traced_per_system = [] +for i, (dataset, positions, plane_index) in enumerate( + zip(dataset_list, positions_list, plane_indices) +): + + def trace_positions(positions_arr, _plane_index=plane_index): + grid = al.Grid2DIrregular(values=positions_arr, xp=jnp) + traced = tracer.traced_grid_2d_list_from(grid=grid, xp=jnp)[_plane_index] + return traced.array if hasattr(traced, "array") else traced + + _, traced = jit_profile( + trace_positions, + f"step1_trace_{dataset.name}", + jnp.array(positions), + ) + traced_per_system.append(traced) + likelihood_steps.append( + (f"1.{i} ray-trace {dataset.name} ({len(positions)} img)", timer.records[-1][1] / 10) + ) + +# --------------------------------------------------------------------------- +# Step 2 — magnification at the observed positions via the tracer Hessian. +# The production fit weights each source-plane residual by |mu| so the +# image-plane positional noise maps into the source plane correctly. +# --------------------------------------------------------------------------- +import autogalaxy as ag # noqa: E402 + +magnifications_per_system = [] +for i, (dataset, positions, plane_index) in enumerate( + zip(dataset_list, positions_list, plane_indices) +): + lens_calc = ag.LensCalc.from_tracer(tracer=tracer, use_multi_plane=True, plane_j=plane_index) + + def magnifications_at(positions_arr, _lens_calc=lens_calc): + # The raw traced array goes straight into the hessian path: wrapping it in + # Grid2DIrregular here would trip __array__ on the tracer (the hessian + # slices grid[:, 0] directly, which is fine on a bare jax.Array). + mags = _lens_calc.magnification_2d_via_hessian_from(grid=positions_arr, xp=jnp) + mags = mags.array if hasattr(mags, "array") else mags + return jnp.abs(mags) + + _, mags = jit_profile( + magnifications_at, + f"step2_magnification_{dataset.name}", + jnp.array(positions), + ) + magnifications_per_system.append(mags) + likelihood_steps.append( + (f"2.{i} magnification (hessian) {dataset.name}", timer.records[-1][1] / 10) + ) + +# --------------------------------------------------------------------------- +# Step 3 — centroid + magnification-weighted chi-squared per system +# (pure jnp math; tiny next to steps 1-2, but the production formula). +# --------------------------------------------------------------------------- +chi_squared_per_system = [] +for i, (dataset, traced, mags, noise, centre) in enumerate( + zip( + dataset_list, traced_per_system, magnifications_per_system, noise_list, TRUTH_SOURCE_CENTRES + ) +): + + def weighted_chi_squared(traced_arr, mags_arr, noise_arr, centre_arr): + distances_sq = jnp.sum((traced_arr - centre_arr) ** 2, axis=1) + return jnp.sum(distances_sq * mags_arr**2 / noise_arr**2) + + _, chi_sq = jit_profile( + weighted_chi_squared, + f"step3_chi2_{dataset.name}", + jnp.array(traced), + jnp.array(mags), + jnp.array(noise), + jnp.array(centre), + ) + chi_squared_per_system.append(chi_sq) + likelihood_steps.append((f"3.{i} weighted chi2 {dataset.name}", timer.records[-1][1] / 10)) + +# --------------------------------------------------------------------------- +# Step 4 — total log-likelihood assembly: -0.5 * (sum chi2 + noise norm). +# --------------------------------------------------------------------------- +noise_norm_terms = jnp.array( + [ + float(np.sum(np.log(2.0 * np.pi * (np.asarray(noise) / np.asarray(mags)) ** 2))) + for noise, mags in zip(noise_list, magnifications_per_system) + ] +) + + +def log_likelihood_total(chi_squareds, noise_norms): + return -0.5 * (jnp.sum(chi_squareds) + jnp.sum(noise_norms)) + + +_, log_likelihood = jit_profile( + log_likelihood_total, + "step4_log_likelihood", + jnp.array(chi_squared_per_system), + noise_norm_terms, +) +likelihood_steps.append(("4 log-likelihood assembly", timer.records[-1][1] / 10)) + + +# =================================================================== +# PART C — Eager reference check (FitPositionsSource, production path) +# =================================================================== +with timer.section("reference_fit_positions_source"): + reference_log_likelihood = sum( + float( + al.FitPositionsSource( + name=dataset.name, + data=dataset.positions, + noise_map=dataset.positions_noise_map, + tracer=tracer, + solver=None, + ).log_likelihood + ) + for dataset in dataset_list + ) + +print(f"\n step-by-step log likelihood: {float(log_likelihood):.8e}") +print(f" FitPositionsSource reference: {reference_log_likelihood:.8e}") +assert np.isclose(float(log_likelihood), reference_log_likelihood, rtol=1e-4), ( + "per-step decomposition does not match the production FitPositionsSource value" +) +print(" MATCH (rtol=1e-4)") + + +# =================================================================== +# PART D — Summary + artifacts +# =================================================================== +import json + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +al_version = al.__version__ + +timer.summary() + +print("\n" + "=" * 70) +print(f"PER-STEP BREAKDOWN SUMMARY — CLUSTER SOURCE-PLANE — v{al_version}") +print("=" * 70) +max_label = max(len(label) for label, _ in likelihood_steps) +step_total = 0.0 +for i, (label, per_call) in enumerate(likelihood_steps, 1): + print(f" {i:>2}. {label:<{max_label}} {per_call:>12.6f} s") + step_total += per_call +print("-" * 70) +print(f" {'TOTAL (step-by-step)':<{max_label}} {step_total:>12.6f} s") +print("=" * 70) + +breakdown_summary = { + "autolens_version": al_version, + "device": device_info_dict(), + "configuration": { + "n_systems": len(dataset_list), + "n_images_total": int(sum(len(p) for p in positions_list)), + "n_planes": len(tracer.planes), + "n_mass_profiles": n_mass_profiles, + "likelihood": "source_plane (FitPositionsSource)", + }, + "steps": {label: per_call for label, per_call in likelihood_steps}, + "total_step_by_step": step_total, + "reference_log_likelihood": reference_log_likelihood, +} + +dict_path, chart_path = resolve_output_paths( + _cli, + default_dir=_workspace_root / "results" / "breakdown" / "cluster", + default_basename=f"source_plane_breakdown_v{al_version}", +) +dict_path.write_text(json.dumps(breakdown_summary, indent=2)) +print(f"\n Results dict saved to: {dict_path}") + +labels = [label for label, _ in likelihood_steps] +times = [per_call for _, per_call in likelihood_steps] +fig, ax = plt.subplots(figsize=(10, 5)) +y_pos = range(len(labels)) +bars = ax.barh(y_pos, times, color="#4C72B0", edgecolor="white", height=0.6) +for bar, t in zip(bars, times): + ax.text( + bar.get_width() + max(times) * 0.01, + bar.get_y() + bar.get_height() / 2, + f"{t:.6f} s", + va="center", + fontsize=9, + ) +ax.set_yticks(y_pos) +ax.set_yticklabels(labels, fontsize=10) +ax.invert_yaxis() +ax.set_xlabel("Time per call (s)", fontsize=11) +fig.suptitle(f"Cluster source-plane likelihood breakdown — v{al_version}", fontsize=12) +fig.tight_layout() +fig.savefig(chart_path, dpi=150) +print(f" Bar chart saved to: {chart_path}") diff --git a/results/breakdown/cluster/image_plane_breakdown_v2026.7.6.649.json b/results/breakdown/cluster/image_plane_breakdown_v2026.7.6.649.json new file mode 100644 index 0000000..b2803d8 --- /dev/null +++ b/results/breakdown/cluster/image_plane_breakdown_v2026.7.6.649.json @@ -0,0 +1,25 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "cpu", + "device": "TFRT_CPU_0" + }, + "configuration": { + "n_systems": 2, + "n_images_total": 6, + "n_planes": 3, + "n_mass_profiles": 13, + "likelihood": "image_plane (FitPositionsImagePairRepeat)", + "solver_grid": "200x200 @ 0.7 arcsec", + "solver_pixel_scale_precision": 0.01 + }, + "steps": { + "1 source centres (back-trace)": 0.008795200000349723, + "2.point_0 PointSolver solve (z=1.0)": 0.41861359999984415, + "2.point_1 PointSolver solve (z=2.0)": 0.3188671666666778, + "3.point_0 FitPositionsImagePairRepeat (fit total)": 2.271588399999928, + "3.point_1 FitPositionsImagePairRepeat (fit total)": 1.9715207999997801 + }, + "total_step_by_step": 4.98938516666658, + "log_likelihood": -31265964.625419706 +} \ No newline at end of file diff --git a/results/breakdown/cluster/image_plane_breakdown_v2026.7.6.649.png b/results/breakdown/cluster/image_plane_breakdown_v2026.7.6.649.png new file mode 100644 index 0000000..9ef1860 Binary files /dev/null and b/results/breakdown/cluster/image_plane_breakdown_v2026.7.6.649.png differ diff --git a/results/breakdown/cluster/source_plane_breakdown_v2026.7.6.649.json b/results/breakdown/cluster/source_plane_breakdown_v2026.7.6.649.json new file mode 100644 index 0000000..afd6b7f --- /dev/null +++ b/results/breakdown/cluster/source_plane_breakdown_v2026.7.6.649.json @@ -0,0 +1,25 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "cpu", + "device": "TFRT_CPU_0" + }, + "configuration": { + "n_systems": 2, + "n_images_total": 6, + "n_planes": 3, + "n_mass_profiles": 13, + "likelihood": "source_plane (FitPositionsSource)" + }, + "steps": { + "1.0 ray-trace point_0 (3 img)": 0.0008368400000108523, + "1.1 ray-trace point_1 (3 img)": 0.0007585400000607479, + "2.0 magnification (hessian) point_0": 0.0007032299999991665, + "2.1 magnification (hessian) point_1": 0.0007097899999280344, + "3.0 weighted chi2 point_0": 2.73399999969115e-05, + "3.1 weighted chi2 point_1": 3.465999998297775e-05, + "4 log-likelihood assembly": 2.303000001120381e-05 + }, + "total_step_by_step": 0.0030934299999898942, + "reference_log_likelihood": -40135921.91112596 +} \ No newline at end of file diff --git a/results/breakdown/cluster/source_plane_breakdown_v2026.7.6.649.png b/results/breakdown/cluster/source_plane_breakdown_v2026.7.6.649.png new file mode 100644 index 0000000..46e5648 Binary files /dev/null and b/results/breakdown/cluster/source_plane_breakdown_v2026.7.6.649.png differ diff --git a/results/simulators/cluster_summary_v2026.7.6.649.json b/results/simulators/cluster_summary_v2026.7.6.649.json new file mode 100644 index 0000000..5433519 --- /dev/null +++ b/results/simulators/cluster_summary_v2026.7.6.649.json @@ -0,0 +1,59 @@ +{ + "autolens_version": "2026.7.6.649", + "type": "cluster", + "configuration": { + "imaging_grid_shape": [ + 1000, + 1000 + ], + "solver_grid_shape": [ + 800, + 800 + ], + "pixel_scales": 0.1, + "n_sources": 2, + "source_redshifts": [ + 1.0, + 2.0 + ], + "n_lens_galaxies": 2, + "host_halo_mass_at_200": 1995262314968882.8, + "over_sample_sub_sizes": [ + 32, + 8, + 2 + ] + }, + "phases": { + "setup_imaging_grid": 0.2845069000004514, + "setup_galaxies": 0.013920999999754713, + "setup_tracer": 2.1300000298651867e-05, + "register_pytrees": 0.6894228000001021, + "solver_build": 0.08422690000043076, + "jitted_solve_src0_lower": 29.02274850000049, + "jitted_solve_src0_compile": 20.82806900000014, + "jitted_solve_src0_first_call": 7.785073900000498, + "jitted_solve_src0_steady_x3": 18.943083899999692, + "jitted_solve_src1_lower": 1.7909546999999293, + "jitted_solve_src1_compile": 15.781297700000323, + "jitted_solve_src1_first_call": 6.645430900000065, + "jitted_solve_src1_steady_x3": 19.616422000000057, + "setup_psf_simulator": 0.000772299999880488, + "via_tracer_from": 113.28854270000011, + "output_fits": 0.1623942000005627, + "output_scaling_galaxies_csv": 0.0005623000006380607, + "output_point_datasets": 0.0029479999993782258, + "output_json": 0.05425109999941924 + }, + "diagnostic": { + "solver_compile_total_s": 36.609366700000464, + "solver_lower_total_s": 30.81370320000042, + "solver_steady_total_s": 38.55950589999975, + "via_tracer_from_s": 113.28854270000011, + "compile_vs_via_tracer_ratio": 0.32 + }, + "positions_found": { + "src0": 3, + "src1": 3 + } +} \ No newline at end of file diff --git a/results/simulators/cluster_summary_v2026.7.6.649.png b/results/simulators/cluster_summary_v2026.7.6.649.png new file mode 100644 index 0000000..a7d55a9 Binary files /dev/null and b/results/simulators/cluster_summary_v2026.7.6.649.png differ diff --git a/simulators/cluster.py b/simulators/cluster.py index 2c626c1..a0173b9 100644 --- a/simulators/cluster.py +++ b/simulators/cluster.py @@ -138,6 +138,27 @@ def jit_profile(func, label, *args, n_repeats=10): (-0.8, 1.2), ] +# Scaling-tier members: mirrors the autolens_workspace cluster simulator truth (the +# reference-anchored Lenstool convention made default in autolens_workspace#238): +# b0_i = b0_ref * (L_i/L_ref)^0.5, rs_i = rs_ref * (L_i/L_ref)^0.5, brightest member anchors. +scaling_galaxies_centres = [ + (5.5, -6.5), + (-7.5, 3.0), + (12.0, -5.0), + (-4.0, -9.0), + (3.0, 13.0), + (-14.0, 4.0), + (15.0, 9.0), + (-9.0, -12.0), + (8.5, 5.5), + (-6.5, 11.0), +] +scaling_galaxies_luminosities = [0.40, 0.32, 0.25, 0.20, 0.16, 0.13, 0.10, 0.08, 0.06, 0.05] +SCALING_B0_REF = 0.12 +SCALING_RS_REF = 10.0 +SCALING_RA = 0.1 +SCALING_EXPONENT = 0.5 + with timer.section("setup_imaging_grid"): imaging_grid = al.Grid2D.uniform(shape_native=(1000, 1000), pixel_scales=0.1) imaging_over_sample = al.util.over_sample.over_sample_size_via_radial_bins_from( @@ -190,8 +211,29 @@ def jit_profile(func, label, *args, n_repeats=10): point = al.ps.Point(centre=centre) source_galaxies.append(al.Galaxy(redshift=src_z, bulge=bulge, **{f"point_{i}": point})) + scaling_galaxies = [] + _lum_ref = max(scaling_galaxies_luminosities) + for centre, luminosity in zip(scaling_galaxies_centres, scaling_galaxies_luminosities): + ratio = luminosity / _lum_ref + scaling_galaxies.append( + al.Galaxy( + redshift=redshift_lens, + bulge=al.lp.SersicSph( + centre=centre, intensity=luminosity, effective_radius=0.8, sersic_index=3.0 + ), + mass=al.mp.dPIEMassSph( + centre=centre, + ra=SCALING_RA, + rs=SCALING_RS_REF * ratio**SCALING_EXPONENT, + b0=SCALING_B0_REF * ratio**SCALING_EXPONENT, + ), + ) + ) + with timer.section("setup_tracer"): - tracer = al.Tracer(galaxies=main_lens_galaxies + [host_halo_galaxy] + source_galaxies) + tracer = al.Tracer( + galaxies=main_lens_galaxies + scaling_galaxies + [host_halo_galaxy] + source_galaxies + ) # === PART 2 — Pytree registration === @@ -338,6 +380,13 @@ def jitted_solve(tracer, source_plane_coordinate): overwrite=True, ) +with timer.section("output_scaling_galaxies_csv"): + al.galaxy_table_to_csv( + centres=scaling_galaxies_centres, + luminosities=scaling_galaxies_luminosities, + file_path=dataset_path / "scaling_galaxies.csv", + ) + with timer.section("output_point_datasets"): # Position noise = 5 mas (HST PSF-centroiding precision), not the imaging pixel scale. position_noise = 0.005