feat: combined strong+weak lensing features example (weak series step 8)

llms-full.txt

@@ -430,6 +430,12 @@ AUTO-GENERATED by PyAutoBuild — do not edit by hand; regenerate with generate.
 
 ## weak
 
+- [Fit: Combined Strong + Weak Lensing](scripts/weak/features/strong_lensing/fit.py): This script fits the combined strong+weak dataset simulated by `simulator.py` in this folder with a single shared `Tracer`: the imaging data via `FitImaging` and the shear catalogue via `FitWeak`.
+  - Contents: Dataset, Tracer, Imaging Fit, Weak Fit, Joint Likelihood, Shear Profile
+- [Modeling: Combined Strong + Weak Lensing](scripts/weak/features/strong_lensing/modeling.py): This script fits the combined dataset of `simulator.py` — an `Imaging` dataset of strongly lensed arcs and a `WeakDataset` of the surrounding shear field — with a **single lens mass model**, using PyAutoFit's factor-graph API to sample the joint likelihood with one non-linear search.
+  - Contents: Dataset, Model, Analysis List, Analysis Factor & Factor Graph, Search & Model-Fit, Result
+- [Simulator: Combined Strong + Weak Lensing](scripts/weak/features/strong_lensing/simulator.py): This script simulates the two faces of the same gravitational lens: an `Imaging` dataset of its strongly lensed arcs, and a `WeakDataset` of the weak shear it imprints on background galaxies at larger radii. Both are generated from **one** `Tracer`, so the datasets share a single true mass distribution — the setup the fit and modeling examples in this folder use to demonstrate joint strong+weak constraints.
+  - Contents: Dataset Paths, Ray Tracing, Imaging Simulation, Weak Simulation, Output
 - [Fit: Weak Lensing](scripts/weak/fit.py): This script shows how to fit a strong-lens mass model to a weak gravitational lensing shear catalogue. Where the `imaging` and `interferometer` workflows fit a 2D image of a lensed source, the weak-lensing workflow fits a set of (gamma_2, gamma_1) shear measurements at the (y, x) positions of background source galaxies — a ``WeakDataset`` produced by the simulator script in `scripts/weak/simulator.py`.
   - Contents: Dataset, Model, Fit, Visualization, Shear Profile, Notes
 - [__Log Likelihood Function: Weak Lensing__](scripts/weak/likelihood_function.py): This script provides a step-by-step guide of the **PyAutoLens** likelihood function for fitting a lens mass model to a weak gravitational lensing shear catalogue (a `WeakDataset`). It is the weak-lensing companion of the guides for the other dataset types (e.g. `scripts/imaging/likelihood_function.py`), following the same style and level of detail.

scripts/weak/features/strong_lensing/fit.py

@@ -0,0 +1,155 @@
+"""
+Fit: Combined Strong + Weak Lensing
+===================================
+
+This script fits the combined strong+weak dataset simulated by `simulator.py` in this folder with a single
+shared `Tracer`: the imaging data via `FitImaging` and the shear catalogue via `FitWeak`.
+
+Because the two datasets are statistically independent measurements of the same mass distribution, their
+joint log likelihood is simply the sum of the two individual log likelihoods — this additivity is all the
+joint modeling in `modeling.py` needs, and this script makes it explicit before a non-linear search is
+involved.
+
+__Contents__
+
+- **Dataset:** Load both datasets (auto-simulating them if missing) and mask the imaging data.
+- **Tracer:** The single mass model shared by both fits.
+- **Imaging Fit:** Fit the strong-lensing image.
+- **Weak Fit:** Fit the shear catalogue.
+- **Joint Likelihood:** The sum that a joint analysis samples.
+- **Shear Profile:** Data vs model in the space cluster weak lensing is usually shown in.
+"""
+
+from autoconf import jax_wrapper  # Sets JAX environment before other imports
+
+# from autoconf import setup_notebook; setup_notebook()
+
+from pathlib import Path
+
+import autolens as al
+import autolens.plot as aplt
+
+"""
+__Dataset__
+
+Load both halves of the combined dataset from `dataset/weak/strong_lensing/`.
+
+__Dataset Auto-Simulation__
+
+If the dataset does not already exist on your system, it will be created by running the corresponding
+simulator script. This ensures that all example scripts can be run without manually simulating data first.
+"""
+dataset_path = Path("dataset") / "weak" / "strong_lensing"
+
+if al.util.dataset.should_simulate(str(dataset_path)):
+    import subprocess
+    import sys
+
+    subprocess.run(
+        [sys.executable, "scripts/weak/features/strong_lensing/simulator.py"],
+        check=True,
+    )
+
+dataset_imaging = al.Imaging.from_fits(
+    data_path=dataset_path / "data.fits",
+    psf_path=dataset_path / "psf.fits",
+    noise_map_path=dataset_path / "noise_map.fits",
+    pixel_scales=0.1,
+)
+
+dataset_weak = al.from_json(file_path=dataset_path / "dataset.json")
+
+"""
+The imaging data is masked with the standard 3.0" circular mask — inside it live the arcs; the shear
+catalogue carries the information outside it, to 10".
+"""
+mask = al.Mask2D.circular(
+    shape_native=dataset_imaging.shape_native,
+    pixel_scales=dataset_imaging.pixel_scales,
+    radius=3.0,
+)
+
+dataset_imaging = dataset_imaging.apply_mask(mask=mask)
+
+"""
+__Tracer__
+
+One tracer fits both datasets. We use the simulator's true parameters, so both fits below are "perfect"
+up to noise — swap any parameter to see both likelihoods respond, which is exactly what a non-linear
+search exploits in `modeling.py`.
+"""
+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.8, angle=45.0),
+    ),
+)
+
+source_galaxy = al.Galaxy(
+    redshift=1.0,
+    bulge=al.lp.SersicCore(
+        centre=(0.05, 0.05),
+        ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=60.0),
+        intensity=4.0,
+        effective_radius=0.1,
+        sersic_index=1.0,
+    ),
+)
+
+tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])
+
+"""
+__Imaging Fit__
+
+The strong-lensing side is the standard `FitImaging` of every imaging example: model image, PSF convolution,
+residuals and a Gaussian likelihood over the masked pixels.
+"""
+fit_imaging = al.FitImaging(dataset=dataset_imaging, tracer=tracer)
+
+print(f"imaging log_likelihood : {fit_imaging.log_likelihood:.2f}")
+
+aplt.subplot_fit_imaging(fit=fit_imaging, output_path=dataset_path, output_format="png")
+
+"""
+__Weak Fit__
+
+The weak-lensing side is the `FitWeak` of `scripts/weak/fit.py`: the tracer's shear field evaluated at the
+catalogue positions, compared to the measured shears over 2N independent components.
+"""
+fit_weak = al.FitWeak(dataset=dataset_weak, tracer=tracer)
+
+print(f"weak log_likelihood    : {fit_weak.log_likelihood:.2f}")
+print(
+    f"weak chi_squared       : {fit_weak.chi_squared:.1f} "
+    f"(expected ~{2 * dataset_weak.n_galaxies} for the true model)"
+)
+
+aplt.subplot_fit_weak(fit=fit_weak, output_path=dataset_path, output_format="png")
+
+"""
+__Joint Likelihood__
+
+The datasets are independent (different galaxies, different noise), so the joint log likelihood of the
+shared tracer is their sum. This one line is the entire statistical content of "combining strong and weak
+lensing" — `modeling.py` wires it into PyAutoFit's factor-graph API so a non-linear search samples it.
+"""
+log_likelihood_joint = fit_imaging.log_likelihood + fit_weak.log_likelihood
+
+print(f"joint log_likelihood   : {log_likelihood_joint:.2f}")
+
+"""
+__Shear Profile__
+
+The tangential shear profile shows where the weak information lives: the model curve (from the same tracer
+fitting the arcs at 1.6") is tested by the binned data points out to 10" — radii the imaging mask never
+sees. The cross component scattering around zero is the standard B-mode systematics null test.
+"""
+aplt.plot_shear_profile(
+    fit_weak,
+    centre=(0.0, 0.0),
+    bins=8,
+    output_path=dataset_path,
+    output_format="png",
+)

scripts/weak/features/strong_lensing/modeling.py

@@ -0,0 +1,219 @@
+"""
+Modeling: Combined Strong + Weak Lensing
+========================================
+
+This script fits the combined dataset of `simulator.py` — an `Imaging` dataset of strongly lensed arcs and a
+`WeakDataset` of the surrounding shear field — with a **single lens mass model**, using PyAutoFit's
+factor-graph API to sample the joint likelihood with one non-linear search.
+
+__Why combine them?__
+
+The two datasets constrain the same mass distribution in complementary regimes:
+
+ - **Strong lensing** (the arcs) pins the Einstein radius, and the mass centre exquisitely — but only
+   *inside* ~1 Einstein radius, and mass-model families that agree there can diverge immediately outside it.
+
+ - **Weak lensing** (the shear catalogue) measures the mass profile and its ellipticity out to many Einstein
+   radii — noisily per galaxy, but with statistical power in the ensemble, and precisely where the strong
+   lensing has none.
+
+Fitting them jointly forces one parametric mass model to satisfy both, the approach of hybrid-Lenstool's
+joint strong+weak cluster reconstructions (Niemiec et al. 2020, who showed sequential fitting biases the
+profile at 2-3 sigma where a joint fit stays within ~1 sigma) and of the stacked strong+weak analysis of the
+Sloan Giant Arcs Survey group-to-cluster lenses (Oguri et al. 2012).
+
+__Contents__
+
+- **Dataset:** Load both datasets (auto-simulating if missing) and mask the imaging data.
+- **Model:** One lens model whose priors are shared by both datasets' analyses.
+- **Analysis List:** An `AnalysisImaging` and an `AnalysisWeak`, one per dataset.
+- **Analysis Factor & Factor Graph:** Combine them so one search samples the joint likelihood.
+- **Search & Model-Fit:** Nautilus over the shared parameter space.
+- **Result:** The joint constraints, and how to read the strong/weak complementarity in them.
+"""
+
+from autoconf import jax_wrapper  # Sets JAX environment before other imports
+
+# from autoconf import setup_notebook; setup_notebook()
+
+from pathlib import Path
+
+import autofit as af
+import autolens as al
+import autolens.plot as aplt
+
+"""
+__Dataset__
+
+Load both halves of the combined dataset, auto-simulating them if missing (the standard pattern of all
+example scripts), and apply the standard 3.0" circular mask to the imaging data. Everything outside that
+mask is the weak catalogue's territory.
+"""
+dataset_name = "strong_lensing"
+dataset_path = Path("dataset") / "weak" / dataset_name
+
+if al.util.dataset.should_simulate(str(dataset_path)):
+    import subprocess
+    import sys
+
+    subprocess.run(
+        [sys.executable, "scripts/weak/features/strong_lensing/simulator.py"],
+        check=True,
+    )
+
+dataset_imaging = al.Imaging.from_fits(
+    data_path=dataset_path / "data.fits",
+    psf_path=dataset_path / "psf.fits",
+    noise_map_path=dataset_path / "noise_map.fits",
+    pixel_scales=0.1,
+)
+
+mask = al.Mask2D.circular(
+    shape_native=dataset_imaging.shape_native,
+    pixel_scales=dataset_imaging.pixel_scales,
+    radius=3.0,
+)
+
+dataset_imaging = dataset_imaging.apply_mask(mask=mask)
+
+dataset_weak = al.from_json(file_path=dataset_path / "dataset.json")
+
+"""
+__Model__
+
+One model serves both datasets:
+
+ - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters] — the component both
+   datasets constrain, through the arcs inside the mask and the shear outside it.
+
+ - The source galaxy's light is a linear `SersicCore` [6 parameters] — only the imaging dataset sees this;
+   the weak catalogue's galaxies are pure shear probes with no model components.
+
+Crucially the model is composed **once**: passing the same model to both analysis factors below means they
+share the same priors and therefore the same parameters — the definition of a joint fit.
+"""
+# Lens:
+
+mass = af.Model(al.mp.Isothermal)
+
+lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)
+
+# Source:
+
+bulge = af.Model(al.lp_linear.SersicCore)
+
+source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)
+
+# Overall Lens Model:
+
+model = af.Collection(galaxies=af.Collection(lens=lens, source=source))
+
+print(model.info)
+
+"""
+__Analysis List__
+
+One analysis object per dataset, exactly as each would be built in its own modeling script.
+
+The imaging analysis runs in NumPy mode here (`use_jax=False`): the factor graph below evaluates its factors
+in a plain Python loop because the weak-lensing analysis is a NumPy calculation, and mixing an eagerly-JAX
+imaging likelihood into that loop gains nothing over NumPy without JIT compilation. On this small masked
+dataset the NumPy likelihood is fast, and the weak likelihood is fractions of a millisecond.
+"""
+analysis_imaging = al.AnalysisImaging(dataset=dataset_imaging, use_jax=False)
+
+analysis_weak = al.AnalysisWeak(dataset=dataset_weak)
+
+"""
+__Analysis Factor__
+
+Each analysis is wrapped in an `AnalysisFactor` paired with the model. Because both factors receive the
+*same* model object, the factor graph recognises every prior as shared — a single 11-dimensional parameter
+space whose likelihood is the sum of the two factors.
+
+(In the `multi` examples each factor gets a slightly different copy of the model, e.g. per-wavelength
+ellipticities; here total sharing is exactly what "one mass distribution, two datasets" means.)
+"""
+analysis_factor_imaging = af.AnalysisFactor(prior_model=model, analysis=analysis_imaging)
+
+analysis_factor_weak = af.AnalysisFactor(prior_model=model, analysis=analysis_weak)
+
+"""
+__Factor Graph__
+
+The factors combine into a `FactorGraphModel`, whose `global_prior_model` is the shared parameter space and
+whose `log_likelihood_function` is the sum over factors — the quantity `fit.py` computed by hand.
+"""
+factor_graph = af.FactorGraphModel(analysis_factor_imaging, analysis_factor_weak)
+
+"""
+__Search & Model-Fit__
+
+Nautilus samples the joint likelihood. The parameter space is simple (N=11, unimodal), so 100 live points
+suffice; expect the fit to take some minutes on an ordinary CPU (the imaging likelihood dominates the cost —
+adding the weak factor is essentially free, which is much of weak lensing's practical appeal).
+"""
+search = af.Nautilus(
+    path_prefix=Path("weak") / "features",
+    name="strong_lensing_joint",
+    unique_tag=dataset_name,
+    n_live=100,
+    iterations_per_quick_update=10000,
+)
+
+print(
+    """
+    The joint strong+weak non-linear search has begun running.
+
+    This Jupyter notebook cell will progress once the search has completed - this could take some minutes!
+    """
+)
+
+result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)
+
+print("The search has finished run - you may now continue the notebook.")
+
+"""
+__Result__
+
+The search returns one result per factor (imaging first, weak second), sharing a single posterior. The
+mass parameters below are constrained by *both* datasets simultaneously.
+
+To see the complementarity in play, compare this joint posterior to a run with the weak factor removed
+(comment it out of the `FactorGraphModel` above): the Einstein radius barely changes — the arcs own it —
+while the constraints on the mass's elliptical components tighten visibly when the shear at large radius is
+included, because ellipticity is exactly what coherent tangential shear across the field measures.
+"""
+result = result_list[0]
+
+print(result.info)
+
+print(result.max_log_likelihood_instance)
+
+"""
+The per-factor maximum-likelihood fits visualize each dataset's view of the shared model.
+"""
+aplt.subplot_fit_imaging(
+    fit=result_list[0].max_log_likelihood_fit,
+    output_path=dataset_path,
+    output_format="png",
+)
+
+aplt.subplot_fit_weak(
+    fit=result_list[1].max_log_likelihood_fit,
+    output_path=dataset_path,
+    output_format="png",
+)
+
+aplt.corner_anesthetic(samples=result.samples)
+
+"""
+__Wrap Up__
+
+This example closed the loop the weak-lensing series builds towards: one mass model, constrained inside the
+Einstein radius by strong lensing and outside it by weak shear, sampled as a single joint likelihood.
+
+The same factor-graph pattern extends directly to the realistic versions of this analysis: cluster-scale
+lenses with many cluster members (see `scripts/cluster`), real shear catalogues (the upcoming real-data
+example in this series), and any other dataset combination (`scripts/multi`).
+"""