From 3860aeb2e6c91fd682817db3bdf553585479fd3a Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Sat, 16 May 2026 14:12:02 +0100
Subject: [PATCH] feat: mirror simulator profiling scripts (Phase 2)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Phase 2 of the autolens_profiling z_feature roadmap. Mirrors the 6
simulator-profiling scripts from autolens_workspace_developer/jax_profiling/
simulators/ into autolens_profiling/simulators/. _developer stays the
source of truth — the originals do not move.

Files mirrored (6 scripts, ~2040 LOC, all from clean origin/main of _developer):
  jax_profiling/simulators/imaging.py        -> simulators/imaging.py
  jax_profiling/simulators/interferometer.py -> simulators/interferometer.py
  jax_profiling/simulators/point_source.py   -> simulators/point_source.py
  jax_profiling/simulators/cluster.py        -> simulators/cluster.py
  jax_profiling/simulators/group.py          -> simulators/group.py
  jax_profiling/simulators/multi.py          -> simulators/multi.py

Path rewrites applied uniformly across all 6 scripts:
- _workspace_root / "jax_profiling" / "dataset" -> _workspace_root / "dataset"
- _workspace_root / "jax_profiling" / "results" / "simulators"
    -> _workspace_root / "results" / "simulators"
- _script_dir.parents[1] -> _script_dir.parents[0]
    (scripts are one level shallower in the new layout —
     simulators/<name>.py vs jax_profiling/simulators/<name>.py)
- Docstring example commands: `python jax_profiling/simulators/<name>.py`
    -> `python simulators/<name>.py`

The artifact filename convention is unchanged:
  results/simulators/<script>_summary_v<al.__version__>.{json,png}

simulators/README.md replaces the Phase 0 placeholder with a section
narrative covering what each simulator profiles, run instructions
(including the dataset_name collision warning for point_source — its
default writes to dataset/point_source/simple/ which holds Phase 1's
tracked likelihood input JSONs), versioned artifact convention, the
headline run-times table placeholder for Phase 4's dashboard, and a
JIT-readiness audit per script.

Smoke status: py_compile + rewrite verification PASSED for all 6 scripts.
Full runtime smoke (each simulator run-to-completion) was intentionally
skipped — running simulators/point_source.py at default dataset_name="simple"
would overwrite Phase 1's tracked point_dataset_positions_only.json and
tracer.json (corrupting the likelihood inputs). Phase 5's
AUTOLENS_PROFILING_SMOKE=1 short-circuit will provide a clean smoke path;
until then the static checks + matching the Phase 1/3 pattern give
sufficient confidence to ship.

F1 lesson applied: copies came from worktree's clean origin/main of
_developer (NOT the canonical, which is dirty with ~36 modified files).

Closes #4.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 simulators/README.md         |  82 +++++-
 simulators/cluster.py        | 492 +++++++++++++++++++++++++++++++++++
 simulators/group.py          | 344 ++++++++++++++++++++++++
 simulators/imaging.py        | 289 ++++++++++++++++++++
 simulators/interferometer.py | 284 ++++++++++++++++++++
 simulators/multi.py          | 310 ++++++++++++++++++++++
 simulators/point_source.py   | 321 +++++++++++++++++++++++
 7 files changed, 2119 insertions(+), 3 deletions(-)
 create mode 100644 simulators/cluster.py
 create mode 100644 simulators/group.py
 create mode 100644 simulators/imaging.py
 create mode 100644 simulators/interferometer.py
 create mode 100644 simulators/multi.py
 create mode 100644 simulators/point_source.py

diff --git a/simulators/README.md b/simulators/README.md
index f99cc3b..b328765 100644
--- a/simulators/README.md
+++ b/simulators/README.md
@@ -1,7 +1,83 @@
 # simulators
 
-Run-time tracking for the PyAutoLens simulators (imaging, interferometer, point-source, cluster, group, multi).
+Run-time profiling for the PyAutoLens simulators — the scripts that produce mock lensed datasets (imaging, interferometer, point-source, cluster-scale, group-scale, multi-wavelength) used as inputs to fitting pipelines. These scripts profile **simulation cost**, not likelihood cost (which lives under [`likelihood/`](../likelihood/README.md)). They time grid construction, ray-tracing, PSF convolution, point-solver work, FITS / JSON output, and whatever else dominates each science case.
 
-Populated by **Phase 2** of the `autolens_profiling` roadmap. Will mirror `autolens_workspace_developer/jax_profiling/simulators/`. JIT support is partial across the simulator suite — placeholder scripts will land for cases where upstream JIT is not yet implemented, linking to the relevant tracking issues.
+## Why profile the simulators?
 
-See the top-level [README](../README.md) for the full phase plan.
+Slow simulators block iteration. The user reported the cluster-scale simulator as the original headliner: JIT-wrapped `solver.solve` was supposed to make it fast, but didn't, so the per-phase breakdown surfaced exactly which step was the bottleneck. Each script here applies the same per-phase timer pattern.
+
+## Scripts
+
+| Script | What it profiles | Distinguishing notes |
+|--------|------------------|----------------------|
+| [`imaging.py`](./imaging.py) | Single-band imaging simulator (HST / Euclid / JWST / AO presets). Grid + over-sampling setup, tracer + galaxies, `tracer.image_2d_from` (eager + JIT), `simulator.via_tracer_from` (NumPy convolution path), output. | Cheapest — galaxy-scale, single band. Useful baseline. |
+| [`interferometer.py`](./interferometer.py) | Visibility-space simulator with synthetic uv-wavelengths (100 baselines inline, no disk file). DFT transformer setup, then `tracer.image_2d_from` eager + JIT. | No PSF-convolution-on-image path; the FFT is in the transformer. |
+| [`point_source.py`](./point_source.py) | Point-source simulator: lens + source tracer, eager `solver.solve`, JIT-compiled `solver.solve` (`xp=jnp`, `remove_infinities=False`), `tracer.time_delays_from`. Internal **eager vs JIT solver assertion** — surfaces the same upstream `PointSolver` drift tracked in [PyAutoLens#514](https://github.com/PyAutoLabs/PyAutoLens/issues/514). | Smallest grid (200×200 @ 0.05″/px). |
+| [`cluster.py`](./cluster.py) | Cluster-scale simulator — the original headliner. 800×800 grid @ 0.1″/px, multiple source galaxies, per-source JIT-compiled `solver.solve` with compile cost reported per source. | Heaviest — multiple JIT compiles, large grid. |
+| [`group.py`](./group.py) | Group-scale imaging simulator. 250×250 grid with adaptive over-sampling at 3 galaxy centres, 1 main lens + 2 extra galaxies + source. | Mid-weight; tests adaptive over-sampling cost. |
+| [`multi.py`](./multi.py) | Multi-wavelength (two-band: g + r) imaging simulator. Per-band grid setup with adaptive over-sampling, two tracers, per-band convolution. | Tests parallel-band orchestration cost. |
+
+## What each script reports
+
+Per-phase wall-time table printed to stdout, plus a structured JSON written to `results/simulators/<script>_summary_v<al.__version__>.json` carrying:
+
+- Phase timings (one entry per `timer.section("...")` block).
+- Key derived timings (e.g. `solver_solve_eager_s`, `via_tracer_from_s`).
+- Configuration metadata (grid shape, pixel scale, model summary).
+- `al.__version__` tag for cross-release comparison.
+
+A matching `.png` bar chart of the per-phase costs lands at the same path. Phase 4's dashboard reads these JSONs to refresh the run-times table in this README.
+
+## Versioned artifacts
+
+Filename convention:
+
+```
+results/simulators/<script>_summary_v<al.__version__>.{json,png}
+```
+
+Old versions are retained alongside new ones so cross-release trends stay visible.
+
+## Headline run-times (populated by Phase 4)
+
+| Script | Dataset preset | CPU | Laptop GPU | A100 |
+|--------|----------------|-----|-----------|------|
+| `imaging.py` | simple (HST-resolution defaults) | _populated_ | _populated_ | _populated_ |
+| `interferometer.py` | simple (synthetic uv) | _populated_ | _populated_ | _populated_ |
+| `point_source.py` | simple | _populated_ | _populated_ | _populated_ |
+| `cluster.py` | simple | _populated_ | _populated_ | _populated_ |
+| `group.py` | simple | _populated_ | _populated_ | _populated_ |
+| `multi.py` | simple (g+r) | _populated_ | _populated_ | _populated_ |
+
+Numbers are the **total wall time** for the simulator run-to-completion (not per-likelihood, since simulators run once-and-done). Phase 4's `scripts/build_readme.py` auto-fills this from the latest `*_summary_v<version>.json` artifacts.
+
+## Running a script
+
+From the repo root:
+
+```bash
+cd autolens_profiling
+python simulators/imaging.py
+```
+
+Each script defaults to `dataset_name = "simple"` and writes the produced dataset to `dataset/<type>/<name>/`.
+
+**Note**: `dataset/point_source/simple/point_dataset_positions_only.json` and `dataset/point_source/simple/tracer.json` are tracked in this repo as Phase 1 likelihood inputs; running `simulators/point_source.py` without changing the default `dataset_name` will overwrite them. Use a non-conflicting `dataset_name` (e.g. `"smoke"`) when smoking these scripts, or commit the regenerated bytes if the regeneration is intentional.
+
+Other `dataset/<type>/{cluster,group,multi}/` and `*/simple/` variants are gitignored, so writing to them does not pollute the repo.
+
+**Codex / sandboxed runs:**
+
+```bash
+NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib python simulators/imaging.py
+```
+
+## JIT readiness
+
+The prompt for Phase 2 flagged that "JAX jitting for these are not fully implemented everywhere so we may have placeholders". Status today:
+
+- `point_source.py` — JIT path for `solver.solve` exercised and asserted against eager. (The upstream drift in this assertion is tracked under [PyAutoLens#514](https://github.com/PyAutoLabs/PyAutoLens/issues/514).)
+- `imaging.py`, `interferometer.py`, `group.py`, `multi.py` — `tracer.image_2d_from` JIT path exercised; `simulator.via_tracer_from` (PSF convolution) is the NumPy path and is the expected bottleneck.
+- `cluster.py` — per-source JIT-compiled `solver.solve` exercised; per-source compile cost reported.
+
+Where the JIT path isn't yet implemented for a given simulator step, the script measures the NumPy path only; future upstream JIT support will be reported in this section as it lands.
diff --git a/simulators/cluster.py b/simulators/cluster.py
new file mode 100644
index 0000000..3b7835d
--- /dev/null
+++ b/simulators/cluster.py
@@ -0,0 +1,492 @@
+"""
+Simulator Profiling: Cluster Scale (Headliner)
+================================================
+
+Profiles `autolens_workspace/scripts/cluster/simulator.py` — the script the
+user reported as slow despite JIT-wrapped `solver.solve`. Pinpoints where time
+goes by timing every major phase separately:
+
+- Pytree registration (register_model + register_instance_pytree)
+- PointSolver construction (800×800 @ 0.1"/px)
+- JIT-compiled solver.solve PER SOURCE (compile cost shown separately for each)
+  so the per-source compile overhead is visible vs. steady-state
+- `simulator.via_tracer_from` on 1000×1000 over-sampled imaging grid
+  (single call — too slow to repeat; this is the suspected numpy bottleneck)
+- FITS + JSON + CSV output
+
+The summary JSON and bar chart make the gap between `solver_compile_total` and
+`via_tracer_from` clearly visible — that's the diagnostic the user needs.
+
+Run from any path:
+    python simulators/cluster.py
+"""
+
+from autoconf import jax_wrapper  # noqa: F401 — must be first
+
+import json
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+import autofit as af
+import autolens as al
+import autolens.plot as aplt
+
+from autofit.jax import register_model as _register_model_pytrees
+from autoarray.abstract_ndarray import register_instance_pytree
+from autolens.lens.tracer import Tracer
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+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()
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[0]
+
+dataset_name = "simple"
+dataset_path = _workspace_root / "dataset" / "cluster" / dataset_name
+dataset_path.mkdir(parents=True, exist_ok=True)
+
+
+# === PART 1 — Setup ===
+
+print("\n--- PART 1: Setup ---")
+
+redshift_lens = 0.5
+source_redshifts = [1.0, 2.0]
+
+main_lens_centres = [
+    (0.0, 0.0),
+    (10.0, 8.0),
+]
+host_halo_centre = (0.0, 0.0)
+source_centres = [
+    (0.3, 0.5),
+    (-0.8, 1.2),
+]
+
+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(
+        grid=imaging_grid,
+        sub_size_list=[32, 8, 2],
+        radial_list=[0.3, 0.6],
+        centre_list=main_lens_centres,
+    )
+    imaging_grid = imaging_grid.apply_over_sampling(
+        over_sample_size=imaging_over_sample
+    )
+
+with timer.section("setup_galaxies"):
+    main_lens_dpie_params = [
+        (8.0, 20.0, 3.0),
+        (5.0, 12.0, 1.2),
+    ]
+    main_lens_sersic_params = [
+        (1.5, 3.0, 4.0),
+        (0.8, 1.5, 3.5),
+    ]
+    main_lens_galaxies = []
+    for centre, (ra, rs, b0), (intensity, effective_radius, sersic_index) in zip(
+        main_lens_centres, main_lens_dpie_params, main_lens_sersic_params
+    ):
+        bulge = al.lp.SersicSph(
+            centre=centre,
+            intensity=intensity,
+            effective_radius=effective_radius,
+            sersic_index=sersic_index,
+        )
+        mass = al.mp.dPIEMassSph(centre=centre, ra=ra, rs=rs, b0=b0)
+        main_lens_galaxies.append(
+            al.Galaxy(redshift=redshift_lens, bulge=bulge, mass=mass)
+        )
+
+    host_halo = al.mp.NFWMCRLudlowSph(
+        centre=host_halo_centre,
+        mass_at_200=10**15.3,
+        redshift_object=redshift_lens,
+        redshift_source=max(source_redshifts),
+    )
+    host_halo_galaxy = al.Galaxy(redshift=redshift_lens, dark=host_halo)
+
+    source_galaxies = []
+    for i, (centre, src_z) in enumerate(zip(source_centres, source_redshifts)):
+        bulge = al.lp.SersicCore(
+            centre=centre,
+            ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0 + 30.0 * i),
+            intensity=2.0,
+            effective_radius=0.3,
+            sersic_index=1.0,
+        )
+        point = al.ps.Point(centre=centre)
+        source_galaxies.append(
+            al.Galaxy(redshift=src_z, bulge=bulge, **{f"point_{i}": point})
+        )
+
+with timer.section("setup_tracer"):
+    tracer = al.Tracer(
+        galaxies=main_lens_galaxies + [host_halo_galaxy] + source_galaxies
+    )
+
+
+# === PART 2 — Pytree registration ===
+
+print("\n--- PART 2: Pytree registration ---")
+
+with timer.section("register_pytrees"):
+    _lens_models = [
+        af.Model(
+            al.Galaxy,
+            redshift=redshift_lens,
+            bulge=af.Model(
+                al.lp.SersicSph,
+                centre=g.bulge.centre,
+                intensity=g.bulge.intensity,
+                effective_radius=g.bulge.effective_radius,
+                sersic_index=g.bulge.sersic_index,
+            ),
+            mass=af.Model(
+                al.mp.dPIEMassSph,
+                centre=g.mass.centre,
+                ra=g.mass.ra,
+                rs=g.mass.rs,
+                b0=g.mass.b0,
+            ),
+        )
+        for g in main_lens_galaxies
+    ]
+    _halo_model = af.Model(
+        al.Galaxy,
+        redshift=redshift_lens,
+        dark=af.Model(
+            al.mp.NFWMCRLudlowSph,
+            centre=host_halo_centre,
+            mass_at_200=host_halo.mass_at_200,
+            redshift_object=redshift_lens,
+            redshift_source=max(source_redshifts),
+        ),
+    )
+    _source_models = [
+        af.Model(
+            al.Galaxy,
+            redshift=src_z,
+            bulge=af.Model(
+                al.lp.SersicCore,
+                centre=src_centre,
+                ell_comps=al.convert.ell_comps_from(
+                    axis_ratio=0.8, angle=60.0 + 30.0 * i
+                ),
+                intensity=2.0,
+                effective_radius=0.3,
+                sersic_index=1.0,
+            ),
+            **{
+                f"point_{i}": af.Model(al.ps.Point, centre=src_centre),
+            },
+        )
+        for i, (src_centre, src_z) in enumerate(zip(source_centres, source_redshifts))
+    ]
+    _registration_model = af.Collection(
+        galaxies=af.Collection(*(_lens_models + [_halo_model] + _source_models))
+    )
+    _register_model_pytrees(_registration_model)
+    register_instance_pytree(Tracer, no_flatten=("cosmology",))
+
+
+# === PART 3 — PointSolver build ===
+
+print("\n--- PART 3: PointSolver build ---")
+
+with timer.section("solver_build"):
+    solver = al.PointSolver.for_grid(
+        grid=al.Grid2D.uniform(shape_native=(800, 800), pixel_scales=0.1),
+        pixel_scale_precision=0.001,
+        magnification_threshold=0.1,
+    )
+
+
+# === PART 4 — JIT solve per source (compile cost visible per-source) ===
+
+print("\n--- PART 4: jitted_solve per source ---")
+
+
+@jax.jit
+def jitted_solve(tracer, source_plane_coordinate):
+    return solver.solve(
+        tracer=tracer,
+        source_plane_coordinate=source_plane_coordinate,
+        xp=jnp,
+        remove_infinities=False,
+    ).array
+
+
+positions_list = []
+
+for i, src_centre in enumerate(source_centres):
+    coord = jnp.asarray(src_centre)
+
+    # Use jit_profile so lower / compile / first / steady are all recorded
+    # independently for each source — this is what makes per-source compile
+    # cost visible.
+    _, raw = jit_profile(
+        lambda c, _tracer=tracer: jitted_solve(_tracer, c),
+        f"jitted_solve_src{i}",
+        coord,
+        n_repeats=3,
+    )
+
+    raw_np = np.asarray(raw)
+    finite_mask = ~(np.isinf(raw_np).any(axis=1) | np.isnan(raw_np).any(axis=1))
+    positions_list.append(al.Grid2DIrregular(raw_np[finite_mask]))
+    print(f"  src{i}: {len(positions_list[-1])} image positions found")
+
+
+# === PART 5 — simulator.via_tracer_from (expected numpy bottleneck) ===
+
+print("\n--- PART 5: simulator.via_tracer_from (1000×1000 over-sampled) ---")
+
+with timer.section("setup_psf_simulator"):
+    psf = al.Convolver.from_gaussian(
+        shape_native=(11, 11), sigma=0.1, pixel_scales=imaging_grid.pixel_scales
+    )
+    simulator = al.SimulatorImaging(
+        exposure_time=300.0,
+        psf=psf,
+        background_sky_level=0.1,
+        add_poisson_noise_to_data=True,
+    )
+
+# Single call only — too slow to repeat on 1000×1000 over-sampled grid
+np.random.seed(1)
+with timer.section("via_tracer_from"):
+    imaging_dataset = simulator.via_tracer_from(tracer=tracer, grid=imaging_grid)
+
+
+# === PART 6 — Outputs (after timing) ===
+
+print("\n--- PART 6: outputs ---")
+
+with timer.section("output_fits"):
+    aplt.fits_imaging(
+        dataset=imaging_dataset,
+        data_path=dataset_path / "data.fits",
+        psf_path=dataset_path / "psf.fits",
+        noise_map_path=dataset_path / "noise_map.fits",
+        overwrite=True,
+    )
+
+with timer.section("output_point_datasets"):
+    # Position noise = 5 mas (HST PSF-centroiding precision), not the imaging pixel scale.
+    position_noise = 0.005
+    dataset_list = []
+    for i, positions in enumerate(positions_list):
+        ds = al.PointDataset(
+            name=f"point_{i}",
+            positions=positions,
+            positions_noise_map=position_noise,
+            redshift=source_redshifts[i],
+        )
+        dataset_list.append(ds)
+        al.output_to_json(obj=ds, file_path=dataset_path / f"point_dataset_{i}.json")
+    al.output_to_csv(
+        datasets=dataset_list,
+        file_path=dataset_path / "point_datasets.csv",
+    )
+
+with timer.section("output_json"):
+    al.output_to_json(obj=tracer, file_path=dataset_path / "tracer.json")
+    al.output_to_json(
+        obj=al.Grid2DIrregular(main_lens_centres),
+        file_path=dataset_path / "main_lens_centres.json",
+    )
+    al.output_to_json(
+        obj=al.Grid2DIrregular([host_halo_centre]),
+        file_path=dataset_path / "host_halo_centre.json",
+    )
+    al.output_to_json(
+        obj=al.Grid2DIrregular(source_centres),
+        file_path=dataset_path / "source_centres.json",
+    )
+
+
+# === Summary — make compile vs via_tracer_from gap obvious ===
+
+al_version = al.__version__
+results_dir = _workspace_root / "results" / "simulators"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+phases = dict(timer.records)
+
+# Aggregate compile totals across both sources
+solver_compile_total = sum(
+    v for k, v in phases.items() if "compile" in k and "solve" in k
+)
+solver_lower_total = sum(v for k, v in phases.items() if "lower" in k and "solve" in k)
+solver_steady_total = sum(
+    v for k, v in phases.items() if "steady" in k and "solve" in k
+)
+
+via_tracer_time = phases.get("via_tracer_from", 0.0)
+
+print("\n" + "=" * 70)
+print("CLUSTER DIAGNOSTIC: compile vs via_tracer_from")
+print("=" * 70)
+print(f"  solver_lower_total  (both sources):  {solver_lower_total:.4f} s")
+print(f"  solver_compile_total (both sources): {solver_compile_total:.4f} s")
+print(f"  solver_steady_total  (both sources): {solver_steady_total:.4f} s")
+print(f"  via_tracer_from:                     {via_tracer_time:.4f} s")
+if via_tracer_time > 0:
+    ratio = (
+        solver_compile_total / via_tracer_time if via_tracer_time > 0 else float("inf")
+    )
+    print(f"  compile / via_tracer ratio:          {ratio:.2f}x")
+print("=" * 70)
+
+results_summary = {
+    "autolens_version": al_version,
+    "type": "cluster",
+    "configuration": {
+        "imaging_grid_shape": [1000, 1000],
+        "solver_grid_shape": [800, 800],
+        "pixel_scales": 0.1,
+        "n_sources": len(source_centres),
+        "source_redshifts": source_redshifts,
+        "n_lens_galaxies": len(main_lens_centres),
+        "host_halo_mass_at_200": 10**15.3,
+        "over_sample_sub_sizes": [32, 8, 2],
+    },
+    "phases": phases,
+    "diagnostic": {
+        "solver_compile_total_s": solver_compile_total,
+        "solver_lower_total_s": solver_lower_total,
+        "solver_steady_total_s": solver_steady_total,
+        "via_tracer_from_s": via_tracer_time,
+        "compile_vs_via_tracer_ratio": (
+            round(solver_compile_total / via_tracer_time, 2)
+            if via_tracer_time > 0
+            else None
+        ),
+    },
+    "positions_found": {f"src{i}": len(p) for i, p in enumerate(positions_list)},
+}
+
+json_path = results_dir / f"cluster_summary_v{al_version}.json"
+json_path.write_text(json.dumps(results_summary, indent=2))
+print(f"\n  Results saved to: {json_path}")
+
+# Bar chart — all phases; compile phases highlighted
+labels = [r[0] for r in timer.records]
+times = [r[1] for r in timer.records]
+
+# Colour compile phases red, via_tracer_from orange, rest blue
+colors = []
+for lbl in labels:
+    if "compile" in lbl:
+        colors.append("#C44E52")
+    elif "via_tracer" in lbl:
+        colors.append("#DD8452")
+    elif "lower" in lbl:
+        colors.append("#937860")
+    else:
+        colors.append("#4C72B0")
+
+fig, ax = plt.subplots(figsize=(13, max(5.0, len(labels) * 0.45)))
+y_pos = range(len(labels))
+bars = ax.barh(y_pos, times, color=colors, 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:.4f} s",
+        va="center",
+        fontsize=8,
+    )
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=8)
+ax.invert_yaxis()
+ax.set_xlabel("Time (s)", fontsize=11)
+fig.suptitle(
+    "Simulator Profiling: Cluster Scale (Headliner)",
+    fontsize=12,
+    fontweight="bold",
+)
+ax.set_title(
+    f"AutoLens v{al_version}  |  1000×1000 imaging / 800×800 solver  |  "
+    f"red=compile, orange=via_tracer",
+    fontsize=8,
+)
+ax.margins(x=0.22)
+fig.tight_layout()
+chart_path = results_dir / f"cluster_summary_v{al_version}.png"
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to: {chart_path}")
+
+timer.summary()
diff --git a/simulators/group.py b/simulators/group.py
new file mode 100644
index 0000000..8d236cd
--- /dev/null
+++ b/simulators/group.py
@@ -0,0 +1,344 @@
+"""
+Simulator Profiling: Group-Scale Imaging
+==========================================
+
+Profiles `autolens_workspace/scripts/group/simulator.py` to pinpoint where
+time goes during simulation. Times the following phases:
+
+- Grid setup (250×250 @ 0.1"/px) with adaptive over-sampling at 3 centres
+- Galaxy construction: 1 main lens + 2 extra galaxies + source
+- Tracer construction
+- `tracer.image_2d_from` (eager + JIT; group-scale grid is large)
+- `simulator.via_tracer_from` (numpy convolution over 250×250)
+- `solver.solve` on 500×500 grid (eager + JIT; group-scale point solving is large)
+- FITS + JSON output (centres, positions, tracer)
+
+Run from any path:
+    python simulators/group.py
+"""
+
+from autoconf import jax_wrapper  # noqa: F401 — must be first
+
+import json
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+import autolens as al
+import autolens.plot as aplt
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+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()
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[0]
+
+dataset_name = "simple"
+dataset_path = _workspace_root / "dataset" / "group" / dataset_name
+dataset_path.mkdir(parents=True, exist_ok=True)
+
+main_lens_centres = [(0.0, 0.0)]
+extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]
+
+
+# === PART 1 — Setup ===
+
+print("\n--- PART 1: Setup ---")
+
+with timer.section("setup_grids"):
+    grid = al.Grid2D.uniform(shape_native=(250, 250), pixel_scales=0.1)
+    over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(
+        grid=grid,
+        sub_size_list=[32, 8, 2],
+        radial_list=[0.3, 0.6],
+        centre_list=main_lens_centres + extra_galaxies_centres,
+    )
+    grid = grid.apply_over_sampling(over_sample_size=over_sample_size)
+
+with timer.section("setup_psf_simulator"):
+    psf = al.Convolver.from_gaussian(
+        shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales
+    )
+    simulator = al.SimulatorImaging(
+        exposure_time=300.0,
+        psf=psf,
+        background_sky_level=0.1,
+        add_poisson_noise_to_data=True,
+    )
+
+with timer.section("setup_galaxies"):
+    lens_0 = al.Galaxy(
+        redshift=0.5,
+        bulge=al.lp.SersicSph(
+            centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0
+        ),
+        mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),
+    )
+    extra_galaxy_0 = al.Galaxy(
+        redshift=0.5,
+        bulge=al.lp.SersicSph(
+            centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0
+        ),
+        mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),
+    )
+    extra_galaxy_1 = al.Galaxy(
+        redshift=0.5,
+        bulge=al.lp.SersicSph(
+            centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0
+        ),
+        mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),
+    )
+    source_galaxy = al.Galaxy(
+        redshift=1.0,
+        bulge=al.lp.SersicCore(
+            centre=(0.0, 0.1),
+            ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),
+            intensity=3.0,
+            effective_radius=0.4,
+            sersic_index=1.0,
+        ),
+    )
+
+with timer.section("setup_tracer"):
+    tracer = al.Tracer(galaxies=[lens_0, extra_galaxy_0, extra_galaxy_1, source_galaxy])
+
+
+# === PART 2 — image_2d_from: eager + JIT ===
+
+print("\n--- PART 2: tracer.image_2d_from (eager + JIT) ---")
+
+with timer.section("image_2d_eager"):
+    image_eager = tracer.image_2d_from(grid=grid)
+
+
+def _image_fn(grid_array):
+    return tracer.image_2d_from(grid=grid, xp=jnp).array
+
+
+jnp_grid = jnp.asarray(grid.array)
+_, image_jit = jit_profile(_image_fn, "image_2d_jit", jnp_grid)
+
+np.testing.assert_allclose(
+    np.asarray(image_eager.array),
+    np.asarray(image_jit),
+    rtol=1e-4,
+    err_msg="group: eager vs JIT image_2d_from mismatch",
+)
+print("  eager ≡ JIT assertion PASSED")
+
+
+# === PART 3 — via_tracer_from ===
+
+print("\n--- PART 3: simulator.via_tracer_from ---")
+
+np.random.seed(1)
+with timer.section("via_tracer_from"):
+    dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)
+
+
+# === PART 4 — solver.solve on 500×500 grid ===
+
+print("\n--- PART 4: solver.solve on 500×500 grid (eager + JIT) ---")
+
+with timer.section("solver_build"):
+    solver_grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.1)
+    solver = al.PointSolver.for_grid(
+        grid=solver_grid, pixel_scale_precision=0.001, magnification_threshold=0.01
+    )
+
+with timer.section("solver_solve_eager"):
+    positions = solver.solve(
+        tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre
+    )
+
+print(f"  Found {len(positions)} image positions (eager)")
+
+
+# Close over `tracer` so it does not cross the JIT boundary — avoids needing
+# pytree registration for a one-tracer profiler.
+@jax.jit
+def jitted_solve(coord):
+    return solver.solve(
+        tracer=tracer,
+        source_plane_coordinate=coord,
+        xp=jnp,
+        remove_infinities=False,
+    ).array
+
+
+src_coord = jnp.asarray(source_galaxy.bulge.centre)
+_, raw_jit = jit_profile(jitted_solve, "solver_jit", src_coord, n_repeats=5)
+
+raw_np = np.asarray(raw_jit)
+finite_mask = ~(np.isinf(raw_np).any(axis=1) | np.isnan(raw_np).any(axis=1))
+positions_jit = al.Grid2DIrregular(raw_np[finite_mask])
+print(f"  Found {len(positions_jit)} image positions (JIT, after inf-strip)")
+
+np.testing.assert_allclose(
+    np.sort(np.asarray(positions), axis=0),
+    np.sort(np.asarray(positions_jit), axis=0),
+    rtol=1e-4,
+    err_msg="group: eager vs JIT solver.solve positions mismatch",
+)
+print("  eager ≡ JIT solver assertion PASSED")
+
+
+# === PART 5 — outputs ===
+
+print("\n--- PART 5: outputs ---")
+
+with timer.section("output_fits"):
+    aplt.fits_imaging(
+        dataset=dataset,
+        data_path=dataset_path / "data.fits",
+        psf_path=dataset_path / "psf.fits",
+        noise_map_path=dataset_path / "noise_map.fits",
+        overwrite=True,
+    )
+
+with timer.section("output_json"):
+    al.output_to_json(obj=tracer, file_path=dataset_path / "tracer.json")
+    al.output_to_json(
+        obj=al.Grid2DIrregular(main_lens_centres),
+        file_path=dataset_path / "main_lens_centres.json",
+    )
+    al.output_to_json(
+        obj=al.Grid2DIrregular(extra_galaxies_centres),
+        file_path=dataset_path / "extra_galaxies_centres.json",
+    )
+    al.output_to_json(obj=positions, file_path=dataset_path / "positions.json")
+
+
+# === Summary ===
+
+al_version = al.__version__
+results_dir = _workspace_root / "results" / "simulators"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+phases = dict(timer.records)
+
+results_summary = {
+    "autolens_version": al_version,
+    "type": "group",
+    "configuration": {
+        "imaging_grid_shape": [250, 250],
+        "solver_grid_shape": [500, 500],
+        "pixel_scales": 0.1,
+        "n_lens_galaxies": 1,
+        "n_extra_galaxies": 2,
+        "over_sample_centres": main_lens_centres + extra_galaxies_centres,
+    },
+    "phases": phases,
+    "key_timings": {
+        "image_2d_eager_s": phases.get("image_2d_eager"),
+        "via_tracer_from_s": phases.get("via_tracer_from"),
+        "solver_solve_eager_s": phases.get("solver_solve_eager"),
+        "n_positions_found": len(positions),
+    },
+}
+
+json_path = results_dir / f"group_summary_v{al_version}.json"
+json_path.write_text(json.dumps(results_summary, indent=2))
+print(f"\n  Results saved to: {json_path}")
+
+labels = [r[0] for r in timer.records]
+times = [r[1] for r in timer.records]
+colors = plt.cm.tab20.colors[: len(labels)]
+
+fig, ax = plt.subplots(figsize=(12, max(4.0, len(labels) * 0.45)))
+y_pos = range(len(labels))
+bars = ax.barh(y_pos, times, color=colors, 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:.4f} s",
+        va="center",
+        fontsize=8,
+    )
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=9)
+ax.invert_yaxis()
+ax.set_xlabel("Time (s)", fontsize=11)
+fig.suptitle("Simulator Profiling: Group Scale", fontsize=12, fontweight="bold")
+ax.set_title(
+    f"AutoLens v{al_version}  |  250×250 imaging / 500×500 solver  |  "
+    f"3 over-sample centres",
+    fontsize=9,
+)
+ax.margins(x=0.22)
+fig.tight_layout()
+chart_path = results_dir / f"group_summary_v{al_version}.png"
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to: {chart_path}")
+
+timer.summary()
diff --git a/simulators/imaging.py b/simulators/imaging.py
new file mode 100644
index 0000000..704854a
--- /dev/null
+++ b/simulators/imaging.py
@@ -0,0 +1,289 @@
+"""
+Simulator Profiling: Imaging
+============================
+
+Profiles `autolens_workspace/scripts/imaging/simulator.py` to pinpoint where
+time goes during simulation. Times the following phases:
+
+- Grid + over-sampling setup
+- Galaxy and tracer construction
+- `tracer.image_2d_from` (eager + JIT comparison)
+- `simulator.via_tracer_from` (numpy convolution path — expected bottleneck)
+- `solver.solve` (eager only; galaxy-scale is fast)
+- FITS + JSON output
+
+Run from the workspace root or any path:
+    python simulators/imaging.py
+"""
+
+from autoconf import jax_wrapper  # noqa: F401 — must be first
+
+import json
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+import autolens as al
+import autolens.plot as aplt
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+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()
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[0]
+
+dataset_name = "simple"
+dataset_path = _workspace_root / "dataset" / "imaging" / dataset_name
+dataset_path.mkdir(parents=True, exist_ok=True)
+
+
+# === PART 1 — Setup ===
+
+print("\n--- PART 1: Setup ---")
+
+with timer.section("setup_grids"):
+    grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.1)
+    over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(
+        grid=grid,
+        sub_size_list=[32, 8, 2],
+        radial_list=[0.3, 0.6],
+        centre_list=[(0.0, 0.0)],
+    )
+    grid = grid.apply_over_sampling(over_sample_size=over_sample_size)
+
+with timer.section("setup_psf_simulator"):
+    psf = al.Convolver.from_gaussian(
+        shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales
+    )
+    simulator = al.SimulatorImaging(
+        exposure_time=300.0,
+        psf=psf,
+        background_sky_level=0.1,
+        add_poisson_noise_to_data=True,
+    )
+
+with timer.section("setup_galaxies"):
+    lens_galaxy = al.Galaxy(
+        redshift=0.5,
+        bulge=al.lp.Sersic(
+            centre=(0.0, 0.0),
+            ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),
+            intensity=2.0,
+            effective_radius=0.6,
+            sersic_index=3.0,
+        ),
+        mass=al.mp.Isothermal(
+            centre=(0.0, 0.0),
+            einstein_radius=1.6,
+            ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),
+        ),
+        shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),
+    )
+    source_galaxy = al.Galaxy(
+        redshift=1.0,
+        bulge=al.lp.SersicCore(
+            centre=(0.0, 0.0),
+            ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),
+            intensity=4.0,
+            effective_radius=0.1,
+            sersic_index=1.0,
+        ),
+    )
+
+with timer.section("setup_tracer"):
+    tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])
+
+
+# === PART 2 — image_2d_from: eager + JIT ===
+
+print("\n--- PART 2: tracer.image_2d_from (eager + JIT) ---")
+
+with timer.section("image_2d_eager"):
+    image_eager = tracer.image_2d_from(grid=grid)
+
+
+def _image_fn(grid_array):
+    return tracer.image_2d_from(grid=grid, xp=jnp).array
+
+
+jnp_grid = jnp.asarray(grid.array)
+_, image_jit = jit_profile(_image_fn, "image_2d_jit", jnp_grid)
+
+np.testing.assert_allclose(
+    np.asarray(image_eager.array),
+    np.asarray(image_jit),
+    rtol=1e-4,
+    err_msg="imaging: eager vs JIT image_2d_from mismatch",
+)
+print("  eager ≡ JIT assertion PASSED")
+
+
+# === PART 3 — via_tracer_from (bottleneck) ===
+
+print("\n--- PART 3: simulator.via_tracer_from ---")
+
+np.random.seed(1)
+with timer.section("via_tracer_from"):
+    dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)
+
+
+# === PART 4 — solver.solve ===
+
+print("\n--- PART 4: solver.solve ---")
+
+with timer.section("solver_build"):
+    solver = al.PointSolver.for_grid(
+        grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1
+    )
+
+with timer.section("solver_solve_eager"):
+    positions = solver.solve(
+        tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre
+    )
+
+
+# === PART 5 — outputs (after timing) ===
+
+print("\n--- PART 5: outputs ---")
+
+with timer.section("output_fits"):
+    aplt.fits_imaging(
+        dataset=dataset,
+        data_path=dataset_path / "data.fits",
+        psf_path=dataset_path / "psf.fits",
+        noise_map_path=dataset_path / "noise_map.fits",
+        overwrite=True,
+    )
+
+with timer.section("output_json"):
+    al.output_to_json(obj=tracer, file_path=dataset_path / "tracer.json")
+    al.output_to_json(obj=positions, file_path=dataset_path / "positions.json")
+
+
+# === Summary ===
+
+al_version = al.__version__
+results_dir = _workspace_root / "results" / "simulators"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+phases = dict(timer.records)
+
+results_summary = {
+    "autolens_version": al_version,
+    "type": "imaging",
+    "configuration": {
+        "grid_shape": [100, 100],
+        "pixel_scales": 0.1,
+        "psf_sigma": 0.1,
+        "over_sampling_sub_sizes": [32, 8, 2],
+    },
+    "phases": phases,
+    "key_timings": {
+        "image_2d_eager_s": phases.get("image_2d_eager"),
+        "via_tracer_from_s": phases.get("via_tracer_from"),
+        "solver_solve_eager_s": phases.get("solver_solve_eager"),
+    },
+}
+
+json_path = results_dir / f"imaging_summary_v{al_version}.json"
+json_path.write_text(json.dumps(results_summary, indent=2))
+print(f"\n  Results saved to: {json_path}")
+
+# Bar chart
+labels = [r[0] for r in timer.records]
+times = [r[1] for r in timer.records]
+colors = plt.cm.tab20.colors[: len(labels)]
+
+fig, ax = plt.subplots(figsize=(12, max(4.0, len(labels) * 0.45)))
+y_pos = range(len(labels))
+bars = ax.barh(y_pos, times, color=colors, 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:.4f} s",
+        va="center",
+        fontsize=8,
+    )
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=9)
+ax.invert_yaxis()
+ax.set_xlabel("Time (s)", fontsize=11)
+fig.suptitle("Simulator Profiling: Imaging", fontsize=12, fontweight="bold")
+ax.set_title(f'AutoLens v{al_version}  |  100×100 @ 0.1"/px', fontsize=9)
+ax.margins(x=0.22)
+fig.tight_layout()
+chart_path = results_dir / f"imaging_summary_v{al_version}.png"
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to: {chart_path}")
+
+timer.summary()
diff --git a/simulators/interferometer.py b/simulators/interferometer.py
new file mode 100644
index 0000000..31e5ca5
--- /dev/null
+++ b/simulators/interferometer.py
@@ -0,0 +1,284 @@
+"""
+Simulator Profiling: Interferometer
+=====================================
+
+Profiles `autolens_workspace/scripts/interferometer/simulator.py` to pinpoint
+where time goes during simulation. Times the following phases:
+
+- Grid setup (256×256 @ 0.1"/px; no over-sampling for interferometry)
+- Synthetic uv_wavelengths construction (100 baselines inline — no disk file)
+- SimulatorInterferometer + TransformerDFT setup
+- Galaxy and tracer construction
+- `tracer.image_2d_from` (eager + JIT)
+- `simulator.via_tracer_from` (DFT to visibilities — expected bottleneck)
+- `solver.solve` (eager)
+- FITS + JSON output
+
+Run from any path:
+    python simulators/interferometer.py
+"""
+
+from autoconf import jax_wrapper  # noqa: F401 — must be first
+
+import json
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+import autolens as al
+import autolens.plot as aplt
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+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()
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[0]
+
+dataset_name = "simple"
+dataset_path = (
+    _workspace_root / "dataset" / "interferometer" / dataset_name
+)
+dataset_path.mkdir(parents=True, exist_ok=True)
+
+
+# === PART 1 — Setup ===
+
+print("\n--- PART 1: Setup ---")
+
+with timer.section("setup_grids"):
+    # Interferometer does not use over-sampling
+    grid = al.Grid2D.uniform(shape_native=(256, 256), pixel_scales=0.1)
+
+with timer.section("setup_uv_wavelengths"):
+    # Inline synthetic baselines — 100 visibilities covering ~10–500 klambda
+    np.random.seed(42)
+    n_vis = 100
+    uv_wavelengths = np.random.uniform(low=10_000, high=500_000, size=(n_vis, 2))
+
+with timer.section("setup_simulator"):
+    simulator = al.SimulatorInterferometer(
+        uv_wavelengths=uv_wavelengths,
+        exposure_time=300.0,
+        noise_sigma=1000.0,
+        transformer_class=al.TransformerDFT,
+    )
+
+with timer.section("setup_galaxies"):
+    lens_galaxy = al.Galaxy(
+        redshift=0.5,
+        mass=al.mp.Isothermal(
+            centre=(0.0, 0.0),
+            einstein_radius=1.6,
+            ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),
+        ),
+        shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),
+    )
+    source_galaxy = al.Galaxy(
+        redshift=1.0,
+        bulge=al.lp.SersicCore(
+            centre=(0.0, 0.0),
+            ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),
+            intensity=0.3,
+            effective_radius=1.0,
+            sersic_index=2.5,
+        ),
+    )
+
+with timer.section("setup_tracer"):
+    tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])
+
+
+# === PART 2 — image_2d_from: eager + JIT ===
+
+print("\n--- PART 2: tracer.image_2d_from (eager + JIT) ---")
+
+with timer.section("image_2d_eager"):
+    image_eager = tracer.image_2d_from(grid=grid)
+
+
+def _image_fn(grid_array):
+    return tracer.image_2d_from(grid=grid, xp=jnp).array
+
+
+jnp_grid = jnp.asarray(grid.array)
+_, image_jit = jit_profile(_image_fn, "image_2d_jit", jnp_grid)
+
+np.testing.assert_allclose(
+    np.asarray(image_eager.array),
+    np.asarray(image_jit),
+    rtol=1e-4,
+    err_msg="interferometer: eager vs JIT image_2d_from mismatch",
+)
+print("  eager ≡ JIT assertion PASSED")
+
+
+# === PART 3 — via_tracer_from (DFT bottleneck) ===
+
+print("\n--- PART 3: simulator.via_tracer_from (DFT) ---")
+
+np.random.seed(1)
+with timer.section("via_tracer_from"):
+    dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)
+
+
+# === PART 4 — solver.solve ===
+
+print("\n--- PART 4: solver.solve ---")
+
+with timer.section("solver_build"):
+    solver = al.PointSolver.for_grid(
+        grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1
+    )
+
+with timer.section("solver_solve_eager"):
+    positions = solver.solve(
+        tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre
+    )
+
+
+# === PART 5 — outputs ===
+
+print("\n--- PART 5: outputs ---")
+
+with timer.section("output_fits"):
+    aplt.fits_interferometer(
+        dataset=dataset,
+        data_path=dataset_path / "data.fits",
+        noise_map_path=dataset_path / "noise_map.fits",
+        uv_wavelengths_path=dataset_path / "uv_wavelengths.fits",
+        overwrite=True,
+    )
+
+with timer.section("output_json"):
+    al.output_to_json(obj=tracer, file_path=dataset_path / "tracer.json")
+
+
+# === Summary ===
+
+al_version = al.__version__
+results_dir = _workspace_root / "results" / "simulators"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+phases = dict(timer.records)
+
+results_summary = {
+    "autolens_version": al_version,
+    "type": "interferometer",
+    "configuration": {
+        "grid_shape": [256, 256],
+        "pixel_scales": 0.1,
+        "n_visibilities": n_vis,
+        "transformer": "TransformerDFT",
+    },
+    "phases": phases,
+    "key_timings": {
+        "image_2d_eager_s": phases.get("image_2d_eager"),
+        "via_tracer_from_s": phases.get("via_tracer_from"),
+        "solver_solve_eager_s": phases.get("solver_solve_eager"),
+    },
+}
+
+json_path = results_dir / f"interferometer_summary_v{al_version}.json"
+json_path.write_text(json.dumps(results_summary, indent=2))
+print(f"\n  Results saved to: {json_path}")
+
+labels = [r[0] for r in timer.records]
+times = [r[1] for r in timer.records]
+colors = plt.cm.tab20.colors[: len(labels)]
+
+fig, ax = plt.subplots(figsize=(12, max(4.0, len(labels) * 0.45)))
+y_pos = range(len(labels))
+bars = ax.barh(y_pos, times, color=colors, 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:.4f} s",
+        va="center",
+        fontsize=8,
+    )
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=9)
+ax.invert_yaxis()
+ax.set_xlabel("Time (s)", fontsize=11)
+fig.suptitle("Simulator Profiling: Interferometer", fontsize=12, fontweight="bold")
+ax.set_title(
+    f'AutoLens v{al_version}  |  256×256 @ 0.1"/px  |  {n_vis} visibilities (DFT)',
+    fontsize=9,
+)
+ax.margins(x=0.22)
+fig.tight_layout()
+chart_path = results_dir / f"interferometer_summary_v{al_version}.png"
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to: {chart_path}")
+
+timer.summary()
diff --git a/simulators/multi.py b/simulators/multi.py
new file mode 100644
index 0000000..7e65ae2
--- /dev/null
+++ b/simulators/multi.py
@@ -0,0 +1,310 @@
+"""
+Simulator Profiling: Multi-Wavelength Imaging
+==============================================
+
+Profiles `autolens_workspace/scripts/multi/simulator.py` to pinpoint where
+time goes when simulating two-waveband imaging (g-band + r-band). Times:
+
+- Per-band grid setup with adaptive over-sampling
+- Per-band PSF and simulator construction
+- Galaxy and tracer construction (two tracers, one per waveband)
+- Per-band `tracer.image_2d_from` (eager + JIT)
+- Per-band `simulator.via_tracer_from` (convolution path)
+- FITS + JSON output per band
+
+Run from any path:
+    python simulators/multi.py
+"""
+
+from autoconf import jax_wrapper  # noqa: F401 — must be first
+
+import json
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+import autolens as al
+import autolens.plot as aplt
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+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()
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[0]
+
+dataset_name = "simple"
+dataset_path = (
+    _workspace_root / "dataset" / "multi" / "imaging" / "lens_sersic"
+)
+dataset_path.mkdir(parents=True, exist_ok=True)
+
+waveband_list = ["g", "r"]
+pixel_scales_list = [0.08, 0.12]
+sigma_list = [0.1, 0.2]
+background_sky_level_list = [0.1, 0.15]
+
+
+# === PART 1 — Setup ===
+
+print("\n--- PART 1: Setup ---")
+
+with timer.section("setup_grids"):
+    grid_list = []
+    for pixel_scales in pixel_scales_list:
+        grid = al.Grid2D.uniform(shape_native=(150, 150), pixel_scales=pixel_scales)
+        over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(
+            grid=grid,
+            sub_size_list=[32, 8, 2],
+            radial_list=[0.3, 0.6],
+            centre_list=[(0.0, 0.0)],
+        )
+        grid_list.append(grid.apply_over_sampling(over_sample_size=over_sample_size))
+
+with timer.section("setup_psf_simulators"):
+    psf_list = [
+        al.Convolver.from_gaussian(
+            shape_native=(11, 11), sigma=sigma, pixel_scales=grid.pixel_scales
+        )
+        for grid, sigma in zip(grid_list, sigma_list)
+    ]
+    simulator_list = [
+        al.SimulatorImaging(
+            exposure_time=300.0,
+            psf=psf,
+            background_sky_level=background_sky_level,
+            add_poisson_noise_to_data=True,
+        )
+        for psf, background_sky_level in zip(psf_list, background_sky_level_list)
+    ]
+
+with timer.section("setup_galaxies"):
+    # Lens intensities differ per band; mass is shared
+    intensity_list_lens = [0.05, 1.5]
+    mass = al.mp.Isothermal(
+        centre=(0.0, 0.0),
+        einstein_radius=1.6,
+        ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),
+    )
+    lens_galaxy_list = [
+        al.Galaxy(
+            redshift=0.5,
+            bulge=al.lp.Sersic(
+                centre=(0.0, 0.0),
+                ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),
+                intensity=intensity,
+                effective_radius=0.8,
+                sersic_index=4.0,
+            ),
+            mass=mass,
+            shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),
+        )
+        for intensity in intensity_list_lens
+    ]
+
+    intensity_list_source = [0.5, 0.7]
+    source_galaxy_list = [
+        al.Galaxy(
+            redshift=1.0,
+            bulge=al.lp.SersicCore(
+                centre=(0.0, 0.0),
+                ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),
+                intensity=intensity,
+                effective_radius=0.1,
+                sersic_index=1.0,
+            ),
+        )
+        for intensity in intensity_list_source
+    ]
+
+with timer.section("setup_tracers"):
+    tracer_list = [
+        al.Tracer(galaxies=[lens_galaxy, source_galaxy])
+        for lens_galaxy, source_galaxy in zip(lens_galaxy_list, source_galaxy_list)
+    ]
+
+
+# === PART 2 — image_2d_from per band: eager + JIT ===
+
+print("\n--- PART 2: tracer.image_2d_from per band (eager + JIT) ---")
+
+for band, tracer, grid in zip(waveband_list, tracer_list, grid_list):
+    with timer.section(f"image_2d_eager_{band}"):
+        image_eager = tracer.image_2d_from(grid=grid)
+
+    _tracer = tracer
+    _grid = grid
+
+    def _image_fn(grid_array, _tracer=_tracer, _grid=_grid):
+        return _tracer.image_2d_from(grid=_grid, xp=jnp).array
+
+    jnp_grid = jnp.asarray(grid.array)
+    _, image_jit = jit_profile(_image_fn, f"image_2d_jit_{band}", jnp_grid)
+
+    np.testing.assert_allclose(
+        np.asarray(image_eager.array),
+        np.asarray(image_jit),
+        rtol=1e-4,
+        err_msg=f"multi/{band}: eager vs JIT image_2d_from mismatch",
+    )
+    print(f"  {band}-band: eager ≡ JIT assertion PASSED")
+
+
+# === PART 3 — via_tracer_from per band ===
+
+print("\n--- PART 3: simulator.via_tracer_from per band ---")
+
+np.random.seed(1)
+dataset_list = []
+for band, simulator, tracer, grid in zip(
+    waveband_list, simulator_list, tracer_list, grid_list
+):
+    with timer.section(f"via_tracer_from_{band}"):
+        dataset_list.append(simulator.via_tracer_from(tracer=tracer, grid=grid))
+
+
+# === PART 4 — outputs ===
+
+print("\n--- PART 4: outputs ---")
+
+with timer.section("output_fits"):
+    for band, dataset in zip(waveband_list, dataset_list):
+        aplt.fits_imaging(
+            dataset=dataset,
+            data_path=dataset_path / f"{band}_data.fits",
+            psf_path=dataset_path / f"{band}_psf.fits",
+            noise_map_path=dataset_path / f"{band}_noise_map.fits",
+            overwrite=True,
+        )
+
+with timer.section("output_json"):
+    for band, tracer in zip(waveband_list, tracer_list):
+        al.output_to_json(obj=tracer, file_path=dataset_path / f"{band}_tracer.json")
+
+
+# === Summary ===
+
+al_version = al.__version__
+results_dir = _workspace_root / "results" / "simulators"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+phases = dict(timer.records)
+
+results_summary = {
+    "autolens_version": al_version,
+    "type": "multi",
+    "configuration": {
+        "bands": waveband_list,
+        "grid_shape": [150, 150],
+        "pixel_scales": pixel_scales_list,
+        "psf_sigmas": sigma_list,
+    },
+    "phases": phases,
+    "key_timings": {
+        "via_tracer_from_g_s": phases.get("via_tracer_from_g"),
+        "via_tracer_from_r_s": phases.get("via_tracer_from_r"),
+    },
+}
+
+json_path = results_dir / f"multi_summary_v{al_version}.json"
+json_path.write_text(json.dumps(results_summary, indent=2))
+print(f"\n  Results saved to: {json_path}")
+
+labels = [r[0] for r in timer.records]
+times = [r[1] for r in timer.records]
+colors = plt.cm.tab20.colors[: len(labels)]
+
+fig, ax = plt.subplots(figsize=(12, max(4.0, len(labels) * 0.45)))
+y_pos = range(len(labels))
+bars = ax.barh(y_pos, times, color=colors, 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:.4f} s",
+        va="center",
+        fontsize=8,
+    )
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=9)
+ax.invert_yaxis()
+ax.set_xlabel("Time (s)", fontsize=11)
+fig.suptitle(
+    "Simulator Profiling: Multi-Wavelength Imaging", fontsize=12, fontweight="bold"
+)
+ax.set_title(f"AutoLens v{al_version}  |  150×150  |  g+r bands", fontsize=9)
+ax.margins(x=0.22)
+fig.tight_layout()
+chart_path = results_dir / f"multi_summary_v{al_version}.png"
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to: {chart_path}")
+
+timer.summary()
diff --git a/simulators/point_source.py b/simulators/point_source.py
new file mode 100644
index 0000000..302b1c8
--- /dev/null
+++ b/simulators/point_source.py
@@ -0,0 +1,321 @@
+"""
+Simulator Profiling: Point Source
+===================================
+
+Profiles `autolens_workspace/scripts/point_source/simulator.py` to pinpoint
+where time goes during simulation. Times the following phases:
+
+- Grid setup (200×200 @ 0.05"/px)
+- Lens + source galaxy and tracer construction
+- `solver.solve` eager
+- JIT-compiled `solver.solve` (xp=jnp, remove_infinities=False)
+- `tracer.time_delays_from` on solved positions
+- `simulator.via_tracer_from` (imaging side — PSF convolution)
+- Point dataset + CSV + FITS + JSON output
+
+Run from any path:
+    python simulators/point_source.py
+"""
+
+from autoconf import jax_wrapper  # noqa: F401 — must be first
+
+import json
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+import autolens as al
+import autolens.plot as aplt
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+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()
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[0]
+
+dataset_name = "simple"
+dataset_path = (
+    _workspace_root / "dataset" / "point_source" / dataset_name
+)
+dataset_path.mkdir(parents=True, exist_ok=True)
+
+
+# === PART 1 — Setup ===
+
+print("\n--- PART 1: Setup ---")
+
+with timer.section("setup_grids"):
+    grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.05)
+
+with timer.section("setup_galaxies"):
+    lens_galaxy = al.Galaxy(
+        redshift=0.5,
+        mass=al.mp.Isothermal(
+            centre=(0.0, 0.0),
+            einstein_radius=1.6,
+            ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),
+        ),
+    )
+    source_galaxy = al.Galaxy(
+        redshift=1.0,
+        light=al.lp.ExponentialCore(
+            centre=(0.07, 0.07),
+            intensity=0.1,
+            effective_radius=0.02,
+            radius_break=0.025,
+        ),
+        point_0=al.ps.Point(centre=(0.07, 0.07)),
+    )
+
+with timer.section("setup_tracer"):
+    tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])
+
+with timer.section("solver_build"):
+    solver = al.PointSolver.for_grid(
+        grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1
+    )
+
+
+# === PART 2 — solver.solve (eager) ===
+
+print("\n--- PART 2: solver.solve (eager) ---")
+
+with timer.section("solver_solve_eager"):
+    positions = solver.solve(
+        tracer=tracer, source_plane_coordinate=source_galaxy.point_0.centre
+    )
+
+print(f"  Found {len(positions)} image positions")
+
+
+# === PART 3 — solver.solve (JIT) ===
+
+print("\n--- PART 3: solver.solve (JIT) ---")
+
+
+# Close over `tracer` so it does not cross the JIT boundary — avoids needing
+# pytree registration for a one-tracer profiler.
+@jax.jit
+def jitted_solve(source_plane_coordinate):
+    return solver.solve(
+        tracer=tracer,
+        source_plane_coordinate=source_plane_coordinate,
+        xp=jnp,
+        remove_infinities=False,
+    ).array
+
+
+src_coord = jnp.asarray(source_galaxy.point_0.centre)
+
+_, raw_jit = jit_profile(jitted_solve, "solver_jit", src_coord, n_repeats=10)
+
+# Strip infinities and compare to eager
+raw_np = np.asarray(raw_jit)
+finite_mask = ~(np.isinf(raw_np).any(axis=1) | np.isnan(raw_np).any(axis=1))
+positions_jit = al.Grid2DIrregular(raw_np[finite_mask])
+
+np.testing.assert_allclose(
+    np.sort(np.asarray(positions), axis=0),
+    np.sort(np.asarray(positions_jit), axis=0),
+    rtol=1e-4,
+    err_msg="point_source: eager vs JIT solver.solve positions mismatch",
+)
+print("  eager ≡ JIT solver assertion PASSED")
+
+
+# === PART 4 — time_delays_from ===
+
+print("\n--- PART 4: tracer.time_delays_from ---")
+
+with timer.section("time_delays_from"):
+    time_delays = tracer.time_delays_from(grid=positions)
+
+
+# === PART 5 — imaging via_tracer_from ===
+
+print("\n--- PART 5: simulator.via_tracer_from (imaging) ---")
+
+with timer.section("setup_psf_simulator"):
+    psf = al.Convolver.from_gaussian(
+        shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales
+    )
+    imaging_simulator = al.SimulatorImaging(
+        exposure_time=300.0,
+        psf=psf,
+        background_sky_level=0.1,
+        add_poisson_noise_to_data=True,
+    )
+
+np.random.seed(1)
+with timer.section("via_tracer_from"):
+    imaging = imaging_simulator.via_tracer_from(tracer=tracer, grid=grid)
+
+
+# === PART 6 — outputs ===
+
+print("\n--- PART 6: outputs ---")
+
+with timer.section("output_point_datasets"):
+    # Position noise = 5 mas (HST PSF-centroiding precision), not the imaging pixel scale.
+    position_noise = 0.005
+    positions_with_noise = positions + np.random.normal(
+        loc=0.0, scale=position_noise, size=positions.shape
+    )
+    positions_with_noise = al.Grid2DIrregular(values=positions_with_noise)
+    dataset = al.PointDataset(
+        name="point_0",
+        positions=positions_with_noise,
+        positions_noise_map=position_noise,
+    )
+    al.output_to_json(
+        obj=dataset,
+        file_path=dataset_path / "point_dataset_positions_only.json",
+    )
+    dataset.to_csv(
+        file_path=dataset_path / "point_dataset_positions_only.csv",
+    )
+
+with timer.section("output_fits"):
+    aplt.fits_imaging(
+        dataset=imaging,
+        data_path=dataset_path / "data.fits",
+        psf_path=dataset_path / "psf.fits",
+        noise_map_path=dataset_path / "noise_map.fits",
+        overwrite=True,
+    )
+
+with timer.section("output_json"):
+    al.output_to_json(obj=tracer, file_path=dataset_path / "tracer.json")
+
+
+# === Summary ===
+
+al_version = al.__version__
+results_dir = _workspace_root / "results" / "simulators"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+phases = dict(timer.records)
+
+results_summary = {
+    "autolens_version": al_version,
+    "type": "point_source",
+    "configuration": {
+        "grid_shape": [200, 200],
+        "pixel_scales": 0.05,
+        "pixel_scale_precision": 0.001,
+        "magnification_threshold": 0.1,
+        "source_centre": list(source_galaxy.point_0.centre),
+    },
+    "phases": phases,
+    "key_timings": {
+        "solver_solve_eager_s": phases.get("solver_solve_eager"),
+        "via_tracer_from_s": phases.get("via_tracer_from"),
+        "time_delays_from_s": phases.get("time_delays_from"),
+        "n_positions_found": len(positions),
+    },
+}
+
+json_path = results_dir / f"point_source_summary_v{al_version}.json"
+json_path.write_text(json.dumps(results_summary, indent=2))
+print(f"\n  Results saved to: {json_path}")
+
+labels = [r[0] for r in timer.records]
+times = [r[1] for r in timer.records]
+colors = plt.cm.tab20.colors[: len(labels)]
+
+fig, ax = plt.subplots(figsize=(12, max(4.0, len(labels) * 0.45)))
+y_pos = range(len(labels))
+bars = ax.barh(y_pos, times, color=colors, 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:.4f} s",
+        va="center",
+        fontsize=8,
+    )
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=9)
+ax.invert_yaxis()
+ax.set_xlabel("Time (s)", fontsize=11)
+fig.suptitle("Simulator Profiling: Point Source", fontsize=12, fontweight="bold")
+ax.set_title(
+    f'AutoLens v{al_version}  |  200×200 @ 0.05"/px  |  {len(positions)} images found',
+    fontsize=9,
+)
+ax.margins(x=0.22)
+fig.tight_layout()
+chart_path = results_dir / f"point_source_summary_v{al_version}.png"
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to: {chart_path}")
+
+timer.summary()